opt/norm: add UseSwapMutation normalization rule

Add a new normalization rule which converts Update and Delete
expressions to UpdateSwap and DeleteSwap, respectively. Crucially, this
normalization rule also replaces the initial Scan of the mutation with a
single-row Values expression, which optimistically provides the complete
old row (if it exists and matches the filters). This allows the
UpdateSwap and DeleteSwap operators to use CPut to validate the
existence of the old row in the same batch as modifying it, changing the
mutation statement from 2 RTs to 1 RT.

With this normalization rule in place, UpdateSwap and DeleteSwap should
be available for use by any qualifying UPDATE and DELETE statements.

Fixes: #144503
Fixes: #85328

Release note (performance improvement): Add new `update swap` and
`delete swap` operators which allow some UPDATE and DELETE statements to
execute in 1 round trip instead of 2 round trips. For now, these
operators can be used when:

- 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 vector indexes modified;
- there are no passthrough columns to RETURNING;
- there are no triggers;
- the table only uses a single column family;
- there are no mutation columns or mutation indexes (i.e. the table is
  not undergoing an ALTER);
- there are no columns using composite encoding (e.g. DECIMAL, FLOAT,
  JSON, etc).

We may lift some of these restrictions in future releases.

Author: michae2

pkg/sql/logictest/testdata/logic_test/swap_mutation

@@ -0,0 +1,127 @@
+# LogicTest: local
+
+# Test that unique indexes work correctly with swap mutations.
+
+statement ok
+CREATE TABLE xyz (
+  x INT PRIMARY KEY,
+  y INT,
+  z INT,
+  UNIQUE INDEX (y),
+  FAMILY (x, y, z)
+)
+
+statement ok
+INSERT INTO xyz VALUES (1, 1, 1), (2, 2, 2)
+
+# This statement should fail with a dupe key error.
+
+statement error duplicate key value violates unique constraint "xyz_y_key"
+UPDATE xyz SET y = 1 WHERE x = 2 AND y = 2 AND z = 2
+
+query III
+SELECT * FROM xyz ORDER BY x
+----
+1  1  1
+2  2  2
+
+# This statement should succeed and modify 0 rows, because no rows have z = 3.
+
+statement ok
+UPDATE xyz SET y = 1 WHERE x = 2 AND y = 2 AND z = 3
+
+query III
+SELECT * FROM xyz ORDER BY x
+----
+1  1  1
+2  2  2
+
+# Test rollback of swap mutations.
+
+statement ok
+BEGIN
+
+statement ok
+UPDATE xyz SET z = 11 WHERE x = 1 AND y = 1 AND z = 1
+
+query III
+SELECT * FROM xyz ORDER BY x
+----
+1  1  11
+2  2  2
+
+statement ok
+SAVEPOINT a
+
+statement ok
+UPDATE xyz SET z = 12 WHERE x = 2 AND y = 2 AND z = 2
+
+query III
+SELECT * FROM xyz ORDER BY x
+----
+1  1  11
+2  2  12
+
+statement ok
+ROLLBACK TO a
+
+statement ok
+UPDATE xyz SET z = 21 WHERE x = 1 AND y = 1 AND z = 11
+
+query III
+SELECT * FROM xyz ORDER BY x
+----
+1  1  21
+2  2  2
+
+statement ok
+SAVEPOINT b
+
+# This should do nothing.
+
+statement ok
+UPDATE xyz SET z = 31 WHERE x = 1 AND y = 1 AND z = 11
+
+query III
+SELECT * FROM xyz ORDER BY x
+----
+1  1  21
+2  2  2
+
+statement ok
+ROLLBACK TO b
+
+query III
+SELECT * FROM xyz ORDER BY x
+----
+1  1  21
+2  2  2
+
+statement ok
+ROLLBACK TO a
+
+query III
+SELECT * FROM xyz ORDER BY x
+----
+1  1  11
+2  2  2
+
+# This should do nothing.
+
+statement ok
+UPDATE xyz SET z = 31 WHERE x = 1 AND y = 1 AND z = 21
+
+query III
+SELECT * FROM xyz ORDER BY x
+----
+1  1  11
+2  2  2
+
+statement ok
+COMMIT
+
+query III
+SELECT * FROM xyz ORDER BY x
+----
+1  1  11
+2  2  2

pkg/sql/logictest/tests/local/generated_test.go

@@ -163,6 +163,25 @@ func (ocl OptionalColList) ToSet() ColSet {
 	return r
 }
 
+// CopyAndMaybeRemapColumns copies this OptionalColList, remapping any columns
+// present in the provided ColMap.
+func (ocl OptionalColList) CopyAndMaybeRemapColumns(m ColMap) (res OptionalColList) {
+	if len(ocl) == 0 {
+		return nil
+	}
+	res = make(OptionalColList, len(ocl))
+	for i, col := range ocl {
+		if col != 0 {
+			if val, ok := m.Get(int(col)); ok {
+				res[i] = ColumnID(val)
+			} else {
+				res[i] = col
+			}
+		}
+	}
+	return res
+}
+
 // ColumnMeta stores information about one of the columns stored in the
 // metadata.
 type ColumnMeta struct {

pkg/sql/opt/column_meta.go

@@ -221,7 +221,7 @@ vectorized: true
 
 # Tracing tests for fast delete.
 statement ok
-CREATE TABLE a (a INT PRIMARY KEY)
+CREATE TABLE a (a DECIMAL PRIMARY KEY)
 
 # Delete range operates in chunks of 600 (defined by row.TableTruncateChunkSize).
 statement ok
@@ -427,28 +427,28 @@ distribution: local
 vectorized: true
 ·
 • render
-│ columns: (a int, b int, "?column?" unknown)
+│ columns: (a decimal, b int, "?column?" unknown)
 │ render ?column?: (NULL)[unknown]
-│ render a: (a)[int]
+│ render a: (a)[decimal]
 │ render b: (b)[int]
 │
 └── • delete
-    │ columns: (a int, b int)
+    │ columns: (a decimal, b int)
     │ estimated row count: 1,000 (missing stats)
     │ from: a
     │ auto commit
     │
     └── • distinct
-        │ columns: (a int, b int)
+        │ columns: (a decimal, b int)
         │ estimated row count: 1,000 (missing stats)
         │ distinct on: a
         │
         └── • cross join (inner)
-            │ columns: (a int, b int)
+            │ columns: (a decimal, b int)
             │ estimated row count: 333,333 (missing stats)
             │
             ├── • scan
-            │     columns: (a int)
+            │     columns: (a decimal)
             │     estimated row count: 1,000 (missing stats)
             │     table: a@a_pkey
             │     spans: FULL SCAN