sql, opt/exec: add UpdateSwap and DeleteSwap to exec factory

Add constructors for UpdateSwap and DeleteSwap to the exec
factory. These closely resemble the constructors for Update and Delete,
but have one additional check:

1. That all columns in the primary index are included in fetch columns.

Informs: #144503

Release note: None

Author: michae2

pkg/sql/delete_swap.go

@@ -142,3 +142,7 @@ func (d *deleteSwapNode) Close(ctx context.Context) {
 func (d *deleteSwapNode) rowsWritten() int64 {
 	return d.run.td.rowsWritten
 }
+
+func (d *deleteSwapNode) enableAutoCommit() {
+	d.run.td.enableAutoCommit()
+}

pkg/sql/distsql_spec_exec_factory.go

@@ -1418,6 +1418,21 @@ func (e *distSQLSpecExecFactory) ConstructUpdate(
 	return nil, unimplemented.NewWithIssue(47473, "experimental opt-driven distsql planning: update")
 }
 
+func (e *distSQLSpecExecFactory) ConstructUpdateSwap(
+	input exec.Node,
+	table cat.Table,
+	fetchCols exec.TableColumnOrdinalSet,
+	updateCols exec.TableColumnOrdinalSet,
+	returnCols exec.TableColumnOrdinalSet,
+	passthrough colinfo.ResultColumns,
+	lockedIndexes cat.IndexOrdinals,
+	autoCommit bool,
+) (exec.Node, error) {
+	return nil, unimplemented.NewWithIssue(
+		47473, "experimental opt-driven distsql planning: update swap",
+	)
+}
+
 func (e *distSQLSpecExecFactory) ConstructUpsert(
 	input exec.Node,
 	table cat.Table,
@@ -1448,6 +1463,20 @@ func (e *distSQLSpecExecFactory) ConstructDelete(
 	return nil, unimplemented.NewWithIssue(47473, "experimental opt-driven distsql planning: delete")
 }
 
+func (e *distSQLSpecExecFactory) ConstructDeleteSwap(
+	input exec.Node,
+	table cat.Table,
+	fetchCols exec.TableColumnOrdinalSet,
+	returnCols exec.TableColumnOrdinalSet,
+	passthrough colinfo.ResultColumns,
+	lockedIndexes cat.IndexOrdinals,
+	autoCommit bool,
+) (exec.Node, error) {
+	return nil, unimplemented.NewWithIssue(
+		47473, "experimental opt-driven distsql planning: delete swap",
+	)
+}
+
 func (e *distSQLSpecExecFactory) ConstructDeleteRange(
 	table cat.Table,
 	needed exec.TableColumnOrdinalSet,

pkg/sql/opt/exec/explain/plan_gist_factory.go

@@ -35,7 +35,7 @@ import (
 )
 
 func init() {
-	if numOperators != 65 {
+	if numOperators != 67 {
 		// This error occurs when an operator has been added or removed in
 		// pkg/sql/opt/exec/explain/factory.opt. If an operator is added at the
 		// end of factory.opt, simply adjust the hardcoded value above. If an

pkg/sql/opt/exec/factory.opt

@@ -652,6 +652,9 @@ define Delete {
 #  - the input to the delete is a scan (without limits);
 #  - there are no inbound FKs to the table.
 #
+# If both DeleteSwap and DeleteRange could be used for a DELETE statement,
+# DeleteSwap takes precedence.
+#
 # See the comment for ConstructScan for descriptions of the needed and
 # indexConstraint parameters, since DeleteRange combines Delete + Scan into a
 # single operator.
@@ -888,3 +891,111 @@ define VectorMutationSearch {
     # the search must return the quantized vector.
     IsIndexPut bool
 }
+
+# UpdateSwap performs an UPDATE statement as a compare-and-swap operation, which
+# only succeeds if the existing row matches what was expected. This
+# compare-and-swap allows the statement to skip the initial fetch from the
+# target table, which changes some simple UPDATE statements from 2 round trips
+# to 1 round trip.
+#
+# This fast path is only possible when certain conditions hold true:
+# - all columns in the primary index are constrained to a single exact value
+#   by the WHERE clause;
+# - only a single row is modified;
+# - there are no FK checks or cascades;
+# - there are no uniqueness checks;
+# - there are no check constraints;
+# - there are no after triggers.
+#
+# Multiple indexes and complex projections are supported.
+#
+# For example, the following statement would be able to use UpdateSwap, as it
+# only updates a single row and all columns are constrained to a single exact
+# value by the WHERE clause.
+#
+#   CREATE TABLE abcd (
+#     a INT PRIMARY KEY,
+#     b INT,
+#     c INT,
+#     d INT AS (b + c) VIRTUAL,
+#     INDEX (b)
+#   );
+#
+#   UPDATE abcd SET a = b + 1, c = 3
+#     WHERE a = 1
+#     AND b IS NOT DISTINCT FROM 2
+#     AND c IS NOT DISTINCT FROM NULL
+#     RETURNING c, d;
+#
+# The parameters are mostly the same as for Update.
+define UpdateSwap {
+    Input exec.Node
+    Table cat.Table
+    FetchCols exec.TableColumnOrdinalSet
+    UpdateCols exec.TableColumnOrdinalSet
+    ReturnCols exec.TableColumnOrdinalSet
+    Passthrough colinfo.ResultColumns
+
+    # If set, the input has already acquired the locks during the initial scan
+    # of the Update (i.e. on the "old" KVs within these indexes).
+    LockedIndexes cat.IndexOrdinals
+
+    # If set, the operator will commit the transaction as part of its execution.
+    AutoCommit bool
+}
+
+# DeleteSwap performs a DELETE statement as a compare-and-swap operation, which
+# only succeeds if the existing row matches what was expected. This
+# compare-and-swap allows the statement to skip the initial fetch from the
+# target table, which changes some simple DELETE statements from 2 round trips
+# to 1 round trip.
+#
+# This fast path is only possible when certain conditions hold true:
+# - all columns in the primary index are constrained to a single exact value
+#   by the WHERE clause;
+# - only a single row is deleted;
+# - there are no FK checks or cascades;
+# - there are no uniqueness checks;
+# - there are no after triggers.
+#
+# Multiple indexes and complex projections are supported.
+#
+# For example, the following statement would be able to use DeleteSwap, as it
+# only deletes a single row and all columns are constrained to a single exact
+# value by the WHERE clause.
+#
+#   CREATE TABLE abcd (
+#     a INT PRIMARY KEY,
+#     b INT,
+#     c INT,
+#     d INT AS (b + c) VIRTUAL,
+#     INDEX (b)
+#   );
+#
+#   DELETE FROM abcd
+#     WHERE a = 1
+#     AND b IS NOT DISTINCT FROM 2
+#     AND c IS NOT DISTINCT FROM NULL
+#     RETURNING *;
+#
+# If both DeleteSwap and DeleteRange could be used for a DELETE statement,
+# DeleteSwap takes precedence.
+#
+# The parameters are the same as for Delete.
+define DeleteSwap {
+    Input exec.Node
+    Table cat.Table
+    FetchCols exec.TableColumnOrdinalSet
+    ReturnCols exec.TableColumnOrdinalSet
+    Passthrough colinfo.ResultColumns
+
+    # If set, the input has already acquired the locks on the specified indexes
+    # on all keys that might be deleted from that index.
+    LockedIndexes cat.IndexOrdinals
+
+    # If set, the operator will commit the transaction as part of its execution.
+    # This is false when executing inside an explicit transaction, or there are
+    # multiple mutations in a statement, or the output of the mutation is
+    # processed through side-effecting expressions.
+    AutoCommit bool
+}

pkg/sql/opt_exec_factory.go

@@ -1630,6 +1630,81 @@ func (ef *execFactory) ConstructUpdate(
 	return &rowCountNode{source: upd}, nil
 }
 
+func (ef *execFactory) ConstructUpdateSwap(
+	input exec.Node,
+	table cat.Table,
+	fetchColOrdSet exec.TableColumnOrdinalSet,
+	updateColOrdSet exec.TableColumnOrdinalSet,
+	returnColOrdSet exec.TableColumnOrdinalSet,
+	passthrough colinfo.ResultColumns,
+	lockedIndexes cat.IndexOrdinals,
+	autoCommit bool,
+) (exec.Node, error) {
+	// TODO(radu): the execution code has an annoying limitation that the fetch
+	// columns must be a superset of the update columns, even when the "old" value
+	// of a column is not necessary. The optimizer code for pruning columns is
+	// aware of this limitation.
+	if !updateColOrdSet.SubsetOf(fetchColOrdSet) {
+		return nil, errors.AssertionFailedf("execution requires all update columns have a fetch column")
+	}
+
+	// For update swap, fetch columns need to include at least every column that
+	// could appear in the primary index.
+	primaryIndex := table.Index(cat.PrimaryIndex)
+	for i := 0; i < primaryIndex.ColumnCount(); i++ {
+		col := primaryIndex.Column(i)
+		if col.Kind() == cat.System {
+			continue
+		}
+		if !fetchColOrdSet.Contains(col.Ordinal()) {
+			return nil, errors.AssertionFailedf("fetch columns missing col %v", col.ColName())
+		}
+	}
+
+	rowsNeeded := !returnColOrdSet.Empty()
+	tabDesc := table.(*optTable).desc
+
+	upd := updateSwapNodePool.Get().(*updateSwapNode)
+	*upd = updateSwapNode{
+		singleInputPlanNode: singleInputPlanNode{input.(planNode)},
+	}
+
+	// If rows are not needed, no columns are returned.
+	var returnCols []catalog.Column
+	if rowsNeeded {
+		returnCols = makeColList(table, returnColOrdSet)
+		upd.columns = colinfo.ResultColumnsFromColumns(tabDesc.GetID(), returnCols)
+		// Add the passthrough columns to the returning columns.
+		upd.columns = append(upd.columns, passthrough...)
+	}
+
+	// Create the table updater, which does the bulk of the work.
+	if err := ef.constructUpdateRun(
+		&upd.run, table, fetchColOrdSet, updateColOrdSet, returnColOrdSet, rowsNeeded,
+		returnCols, exec.CheckOrdinalSet{} /* checks */, passthrough,
+		nil /* uniqueWithTombstoneIndexes */, lockedIndexes,
+	); err != nil {
+		return nil, err
+	}
+
+	if autoCommit {
+		upd.enableAutoCommit()
+	}
+
+	// Serialize the data-modifying plan to ensure that no data is observed that
+	// hasn't been validated first. See the comments on BatchedNext() in
+	// plan_batch.go.
+	if rowsNeeded {
+		return &spoolNode{
+			singleInputPlanNode: singleInputPlanNode{&serializeNode{source: upd}},
+		}, nil
+	}
+
+	// We could use serializeNode here, but using rowCountNode is an
+	// optimization that saves on calls to Next() by the caller.
+	return &rowCountNode{source: upd}, nil
+}
+
 func (ef *execFactory) constructUpdateRun(
 	run *updateRun,
 	table cat.Table,
@@ -1839,6 +1914,71 @@ func (ef *execFactory) ConstructDelete(
 	return &rowCountNode{source: del}, nil
 }
 
+func (ef *execFactory) ConstructDeleteSwap(
+	input exec.Node,
+	table cat.Table,
+	fetchColOrdSet exec.TableColumnOrdinalSet,
+	returnColOrdSet exec.TableColumnOrdinalSet,
+	passthrough colinfo.ResultColumns,
+	lockedIndexes cat.IndexOrdinals,
+	autoCommit bool,
+) (exec.Node, error) {
+	// For delete swap, fetch columns need to include at least every column that
+	// could appear in the primary index.
+	primaryIndex := table.Index(cat.PrimaryIndex)
+	for i := 0; i < primaryIndex.ColumnCount(); i++ {
+		col := primaryIndex.Column(i)
+		if col.Kind() == cat.System {
+			continue
+		}
+		if !fetchColOrdSet.Contains(col.Ordinal()) {
+			return nil, errors.AssertionFailedf("fetch columns missing col %v", col.ColName())
+		}
+	}
+
+	rowsNeeded := !returnColOrdSet.Empty()
+	tabDesc := table.(*optTable).desc
+
+	// Now make a delete node. We use a pool.
+	del := deleteSwapNodePool.Get().(*deleteSwapNode)
+	*del = deleteSwapNode{
+		singleInputPlanNode: singleInputPlanNode{input.(planNode)},
+	}
+
+	// If rows are not needed, no columns are returned.
+	var returnCols []catalog.Column
+	if rowsNeeded {
+		returnCols = makeColList(table, returnColOrdSet)
+		// Delete returns the non-mutation columns specified, in the same
+		// order they are defined in the table.
+		del.columns = colinfo.ResultColumnsFromColumns(tabDesc.GetID(), returnCols)
+		// Add the passthrough columns to the returning columns.
+		del.columns = append(del.columns, passthrough...)
+	}
+
+	// Create the table deleter, which does the bulk of the work.
+	ef.constructDeleteRun(
+		&del.run, table, fetchColOrdSet, rowsNeeded, returnCols, passthrough, lockedIndexes,
+	)
+
+	if autoCommit {
+		del.enableAutoCommit()
+	}
+
+	// Serialize the data-modifying plan to ensure that no data is observed that
+	// hasn't been validated first. See the comments on BatchedNext() in
+	// plan_batch.go.
+	if rowsNeeded {
+		return &spoolNode{
+			singleInputPlanNode: singleInputPlanNode{&serializeNode{source: del}},
+		}, nil
+	}
+
+	// We could use serializeNode here, but using rowCountNode is an
+	// optimization that saves on calls to Next() by the caller.
+	return &rowCountNode{source: del}, nil
+}
+
 func (ef *execFactory) constructDeleteRun(
 	run *deleteRun,
 	table cat.Table,

pkg/sql/update_swap.go

@@ -142,3 +142,7 @@ func (u *updateSwapNode) Close(ctx context.Context) {
 func (u *updateSwapNode) rowsWritten() int64 {
 	return u.run.tu.rowsWritten
 }
+
+func (u *updateSwapNode) enableAutoCommit() {
+	u.run.tu.enableAutoCommit()
+}