opt: support EXPERIMENTAL_RELOCATE

Also fix the formatting of details for split/unsplit.

Release note: None

pkg/sql/opt/bench/stub_factory.go

@@ -305,3 +305,9 @@ func (f *stubFactory) ConstructAlterTableUnsplit(
 func (f *stubFactory) ConstructAlterTableUnsplitAll(index cat.Index) (exec.Node, error) {
 	return struct{}{}, nil
 }
+
+func (f *stubFactory) ConstructAlterTableRelocate(
+	index cat.Index, input exec.Node, relocateLease bool,
+) (exec.Node, error) {
+	return struct{}{}, nil
+}

pkg/sql/opt/exec/execbuilder/relational.go

@@ -253,6 +253,9 @@ func (b *Builder) buildRelational(e memo.RelExpr) (execPlan, error) {
 	case *memo.AlterTableUnsplitAllExpr:
 		ep, err = b.buildAlterTableUnsplitAll(t)
 
+	case *memo.AlterTableRelocateExpr:
+		ep, err = b.buildAlterTableRelocate(t)
+
 	default:
 		if opt.IsSetOp(e) {
 			ep, err = b.buildSetOp(e)

pkg/sql/opt/exec/execbuilder/statement.go

@@ -176,3 +176,24 @@ func (b *Builder) buildAlterTableUnsplitAll(
 	}
 	return ep, nil
 }
+
+func (b *Builder) buildAlterTableRelocate(relocate *memo.AlterTableRelocateExpr) (execPlan, error) {
+	input, err := b.buildRelational(relocate.Input)
+	if err != nil {
+		return execPlan{}, err
+	}
+	table := b.mem.Metadata().Table(relocate.Table)
+	node, err := b.factory.ConstructAlterTableRelocate(
+		table.Index(relocate.Index),
+		input.root,
+		relocate.RelocateLease,
+	)
+	if err != nil {
+		return execPlan{}, err
+	}
+	ep := execPlan{root: node}
+	for i, c := range relocate.Columns {
+		ep.outputCols.Set(int(c), i)
+	}
+	return ep, nil
+}

pkg/sql/opt/exec/factory.go

@@ -416,6 +416,10 @@ type Factory interface {
 	// ConstructAlterTableUnsplitAll creates a node that implements ALTER TABLE/INDEX
 	// UNSPLIT ALL.
 	ConstructAlterTableUnsplitAll(index cat.Index) (Node, error)
+
+	// ConstructAlterTableRelocate creates a node that implements ALTER TABLE/INDEX
+	// UNSPLIT AT.
+	ConstructAlterTableRelocate(index cat.Index, input Node, relocateLease bool) (Node, error)
 }
 
 // OutputOrdering indicates the required output ordering on a Node that is being

pkg/sql/opt/memo/expr_format.go

@@ -187,7 +187,8 @@ func (f *ExprFmtCtx) formatRelational(e RelExpr, tp treeprinter.Node) {
 
 	case *ScanExpr, *VirtualScanExpr, *IndexJoinExpr, *ShowTraceForSessionExpr,
 		*InsertExpr, *UpdateExpr, *UpsertExpr, *DeleteExpr, *SequenceSelectExpr,
-		*WindowExpr, *OpaqueRelExpr:
+		*WindowExpr, *OpaqueRelExpr, *AlterTableSplitExpr, *AlterTableUnsplitExpr,
+		*AlterTableUnsplitAllExpr, *AlterTableRelocateExpr:
 		fmt.Fprintf(f.Buffer, "%v", e.Op())
 		FormatPrivate(f, e.Private(), required)
 
@@ -981,6 +982,12 @@ func FormatPrivate(f *ExprFmtCtx, private interface{}, physProps *physical.Requi
 			fmt.Fprintf(f.Buffer, " %s@%s", tableAlias(f, t.Table), tab.Index(t.Index).Name())
 		}
 
+	case *AlterTableRelocatePrivate:
+		FormatPrivate(f, &t.AlterTableSplitPrivate, nil)
+		if t.RelocateLease {
+			f.Buffer.WriteString(" [lease]")
+		}
+
 	case *JoinPrivate:
 		// Nothing to show; flags are shown separately.
 

pkg/sql/opt/memo/logical_props_builder.go

@@ -755,6 +755,14 @@ func (b *logicalPropsBuilder) buildAlterTableUnsplitAllProps(
 	rel.CanMutate = true
 }
 
+func (b *logicalPropsBuilder) buildAlterTableRelocateProps(
+	relocate *AlterTableRelocateExpr, rel *props.Relational,
+) {
+	b.buildBasicProps(relocate, relocate.Columns, rel)
+	rel.CanHaveSideEffects = true
+	rel.CanMutate = true
+}
+
 func (b *logicalPropsBuilder) buildLimitProps(limit *LimitExpr, rel *props.Relational) {
 	BuildSharedProps(b.mem, limit, &rel.Shared)
 

pkg/sql/opt/ops/statement.opt

@@ -138,3 +138,20 @@ define AlterTableUnsplit {
 define AlterTableUnsplitAll {
     _ AlterTableSplitPrivate
 }
+
+# AlterTableRelocate represents an `ALTER TABLE/INDEX .. SPLIT AT ..` statement.
+[Relational, DDL]
+define AlterTableRelocate {
+    # The input expression provides values for the index columns (or a prefix of
+    # them).
+    Input RelExpr
+
+    _ AlterTableRelocatePrivate
+}
+
+[Private]
+define AlterTableRelocatePrivate {
+    RelocateLease bool
+
+    _ AlterTableSplitPrivate
+}

pkg/sql/opt/optbuilder/alter_table.go

@@ -40,8 +40,8 @@ func (b *Builder) buildAlterTableSplit(split *tree.Split, inScope *scope) (outSc
 
 	b.DisableMemoReuse = true
 
-	// Calculate the desired types for the select statement. It is OK if the
-	// select statement returns fewer columns (the relevant prefix is used).
+	// Calculate the desired types for the input expression. It is OK if it
+	// returns fewer columns (the relevant prefix is used).
 	colNames, colTypes := getIndexColumnNamesAndTypes(index)
 
 	// We don't allow the input statement to reference outer columns, so we
@@ -112,8 +112,8 @@ func (b *Builder) buildAlterTableUnsplit(unsplit *tree.Unsplit, inScope *scope)
 		return outScope
 	}
 
-	// Calculate the desired types for the select statement. It is OK if the
-	// select statement returns fewer columns (the relevant prefix is used).
+	// Calculate the desired types for the input expression. It is OK if it
+	// returns fewer columns (the relevant prefix is used).
 	colNames, colTypes := getIndexColumnNamesAndTypes(index)
 
 	// We don't allow the input statement to reference outer columns, so we
@@ -129,6 +129,64 @@ func (b *Builder) buildAlterTableUnsplit(unsplit *tree.Unsplit, inScope *scope)
 	return outScope
 }
 
+// buildAlterTableRelocate builds an ALTER TABLE/INDEX .. UNSPLIT AT/ALL .. statement.
+func (b *Builder) buildAlterTableRelocate(
+	relocate *tree.Relocate, inScope *scope,
+) (outScope *scope) {
+	flags := cat.Flags{
+		AvoidDescriptorCaches: true,
+		NoTableStats:          true,
+	}
+	index, err := cat.ResolveTableIndex(b.ctx, b.catalog, flags, &relocate.TableOrIndex)
+	if err != nil {
+		panic(builderError{err})
+	}
+	table := index.Table()
+	if err := b.catalog.CheckPrivilege(b.ctx, table, privilege.INSERT); err != nil {
+		panic(builderError{err})
+	}
+
+	b.DisableMemoReuse = true
+
+	outScope = inScope.push()
+	b.synthesizeResultColumns(outScope, sqlbase.AlterTableRelocateColumns)
+
+	// Calculate the desired types for the input expression. It is OK if it
+	// returns fewer columns (the relevant prefix is used).
+	colNames, colTypes := getIndexColumnNamesAndTypes(index)
+
+	// The first column is the target leaseholder or the relocation array,
+	// depending on variant.
+	cmdName := "EXPERIMENTAL_RELOCATE"
+	if relocate.RelocateLease {
+		cmdName += " LEASE"
+		colNames = append([]string{"target leaseholder"}, colNames...)
+		colTypes = append([]*types.T{types.Int}, colTypes...)
+	} else {
+		colNames = append([]string{"relocation array"}, colNames...)
+		colTypes = append([]*types.T{types.IntArray}, colTypes...)
+	}
+
+	// We don't allow the input statement to reference outer columns, so we
+	// pass a "blank" scope rather than inScope.
+	inputScope := b.buildStmt(relocate.Rows, colTypes, &scope{builder: b})
+	checkInputColumns(cmdName, inputScope, colNames, colTypes, 2)
+
+	outScope.expr = b.factory.ConstructAlterTableRelocate(
+		inputScope.expr.(memo.RelExpr),
+		&memo.AlterTableRelocatePrivate{
+			RelocateLease: relocate.RelocateLease,
+			AlterTableSplitPrivate: memo.AlterTableSplitPrivate{
+				Table:   b.factory.Metadata().AddTable(table),
+				Index:   index.Ordinal(),
+				Columns: colsToColList(outScope.cols),
+				Props:   inputScope.makePhysicalProps(),
+			},
+		},
+	)
+	return outScope
+}
+
 // getIndexColumnNamesAndTypes returns the names and types of the index columns.
 func getIndexColumnNamesAndTypes(index cat.Index) (colNames []string, colTypes []*types.T) {
 	colNames = make([]string, index.LaxKeyColumnCount())
@@ -151,7 +209,7 @@ func checkInputColumns(
 		if len(inputScope.cols) == 0 {
 			panic(pgerror.Newf(pgcode.Syntax, "no columns in %s data", context))
 		}
-		panic(pgerror.Newf(pgcode.Syntax, "too few columns in %s data", context))
+		panic(pgerror.Newf(pgcode.Syntax, "less than %d columns in %s data", minPrefix, context))
 	}
 	if len(inputScope.cols) > len(colTypes) {
 		panic(pgerror.Newf(pgcode.Syntax, "too many columns in %s data", context))

pkg/sql/opt/optbuilder/builder.go

@@ -243,6 +243,9 @@ func (b *Builder) buildStmt(
 	case *tree.Unsplit:
 		return b.buildAlterTableUnsplit(stmt, inScope)
 
+	case *tree.Relocate:
+		return b.buildAlterTableRelocate(stmt, inScope)
+
 	default:
 		// See if this statement can be rewritten to another statement using the
 		// delegate functionality.