execbuild: enable lock optimizations for multi-family tables

In #114401 we added some simple optimizations to push locking down into
input operations for SELECT FOR UPDATE under Read Committed
isolation. In #116170 we disabled these optimizations for
multi-column-family tables because we weren't sure that all necessary
column families would be locked.

This commit enables the optimizations for multi-column-family tables,
after confirming that all families needing to be locked are fetched by
the input operation.

Informs: #114566, #116838

Release note: None

Author: michae2

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

@@ -1337,37 +1337,47 @@ func (b *Builder) buildLock(lock *memo.LockExpr) (_ execPlan, outputCols colOrdM
 
 	// In some simple cases we can push the locking down into the input operation
 	// instead of adding what would be a redundant lookup join. We only apply
-	// these optimizations for single-column-family tables.
+	// these optimizations when all necessary column families would be locked by
+	// the input operation.
 	//
 	// TODO(michae2): To optimize more complex cases, such as a Project on top of
 	// a Scan, or multiple Lock ops, etc, we need to do something similar to
 	// ordering. That is, make locking a physical property required by Lock, and
 	// make various operators provide the physical property, with a
 	// locking-semi-LookupJoin as the enforcer of last resort.
 	tab := b.mem.Metadata().Table(lock.Table)
-	if tab.FamilyCount() == 1 {
-		switch input := lock.Input.(type) {
-		case *memo.ScanExpr:
-			if input.Table == lock.KeySource && input.Index == cat.PrimaryIndex {
-				// Make a shallow copy of the scan to avoid mutating the original.
-				scan := *input
-				scan.Locking = scan.Locking.Max(locking)
-				return b.buildRelational(&scan)
-			}
-		case *memo.IndexJoinExpr:
-			if input.Table == lock.KeySource {
-				// Make a shallow copy of the join to avoid mutating the original.
-				join := *input
-				join.Locking = join.Locking.Max(locking)
-				return b.buildRelational(&join)
-			}
-		case *memo.LookupJoinExpr:
-			if input.Table == lock.KeySource && input.Index == cat.PrimaryIndex &&
-				(input.JoinType == opt.InnerJoinOp || input.JoinType == opt.SemiJoinOp) &&
-				!input.IsFirstJoinInPairedJoiner && !input.IsSecondJoinInPairedJoiner &&
-				// We con't push the locking down if the lookup join has additional on
-				// predicates that will filter out rows after the join.
-				len(input.On) == 0 {
+	srcTab := b.mem.Metadata().Table(lock.KeySource)
+	lockFamilies := opt.FamiliesForCols(tab, lock.Table, lock.LockCols)
+	switch input := lock.Input.(type) {
+	case *memo.ScanExpr:
+		if input.Table == lock.KeySource && input.Index == cat.PrimaryIndex &&
+			// Check that all necessary column families would be locked.
+			colFamiliesContainLockFamilies(srcTab, lock.KeySource, input.Cols, lockFamilies) {
+			// Make a shallow copy of the scan to avoid mutating the original.
+			scan := *input
+			scan.Locking = scan.Locking.Max(locking)
+			return b.buildRelational(&scan)
+		}
+	case *memo.IndexJoinExpr:
+		if input.Table == lock.KeySource &&
+			// Check that all necessary column families would be locked.
+			colFamiliesContainLockFamilies(srcTab, lock.KeySource, input.Cols, lockFamilies) {
+			// Make a shallow copy of the join to avoid mutating the original.
+			join := *input
+			join.Locking = join.Locking.Max(locking)
+			return b.buildRelational(&join)
+		}
+	case *memo.LookupJoinExpr:
+		if input.Table == lock.KeySource && input.Index == cat.PrimaryIndex &&
+			(input.JoinType == opt.InnerJoinOp || input.JoinType == opt.SemiJoinOp) &&
+			!input.IsFirstJoinInPairedJoiner && !input.IsSecondJoinInPairedJoiner &&
+			// We con't push the locking down if the lookup join has additional on
+			// predicates that will filter out rows after the join.
+			len(input.On) == 0 {
+			// Check that all necessary column families would be locked.
+			joinInputCols := input.Input.Relational().OutputCols
+			lookupCols := input.Cols.Difference(joinInputCols)
+			if colFamiliesContainLockFamilies(srcTab, lock.KeySource, lookupCols, lockFamilies) {
 				// Make a shallow copy of the join to avoid mutating the original.
 				join := *input
 				join.Locking = join.Locking.Max(locking)
@@ -1393,6 +1403,18 @@ func (b *Builder) buildLock(lock *memo.LockExpr) (_ execPlan, outputCols colOrdM
 	return b.buildLookupJoin(join)
 }
 
+// colFamiliesContainLockFamilies checks that the families for cols include all
+// of the lock families.
+func colFamiliesContainLockFamilies(
+	tab cat.Table, tabID opt.TableID, cols opt.ColSet, lockFamilies intsets.Fast,
+) bool {
+	if tab.FamilyCount() == 1 {
+		return true
+	}
+	families := opt.FamiliesForCols(tab, tabID, cols)
+	return lockFamilies.SubsetOf(families)
+}
+
 func (b *Builder) setMutationFlags(e memo.RelExpr) {
 	b.flags.Set(exec.PlanFlagContainsMutation)
 	switch e.Op() {

pkg/sql/opt/exec/execbuilder/testdata/select_for_update

@@ -2929,7 +2929,8 @@ vectorized: true
       table: t@t_pkey
       spans: /5/0
 
-# Table with multiple column families. Cannot push down locking.
+# Table with multiple column families. Test that we can push down locking in
+# some cases.
 statement ok
 CREATE TABLE xyz (
   x INT PRIMARY KEY,
@@ -2940,6 +2941,7 @@ CREATE TABLE xyz (
   FAMILY (z)
 )
 
+# We can't push down locking if the scan doesn't fetch all families.
 query T
 EXPLAIN (VERBOSE) SELECT y FROM xyz WHERE x = 1 FOR UPDATE
 ----
@@ -2963,48 +2965,93 @@ vectorized: true
           table: xyz@xyz_pkey
           spans: /1/0
 
+# We can push down locking if the scan fetches all families.
+query T
+EXPLAIN (VERBOSE) SELECT y, z FROM xyz WHERE x = 1 FOR UPDATE
+----
+distribution: local
+vectorized: true
+·
+• project
+│ columns: (y, z)
+│
+└── • scan
+      columns: (x, y, z)
+      estimated row count: 1 (missing stats)
+      table: xyz@xyz_pkey
+      spans: /1-/2
+      locking strength: for update
+
+# We can't push down locking if the index join doesn't fetch all families.
+query T
+EXPLAIN (VERBOSE) SELECT y FROM xyz WHERE z = 2 FOR SHARE
+----
+distribution: local
+vectorized: true
+·
+• project
+│ columns: (y)
+│
+└── • lookup join (semi)
+    │ columns: (x, y)
+    │ estimated row count: 10 (missing stats)
+    │ table: xyz@xyz_pkey
+    │ equality: (x) = (x)
+    │ equality cols are key
+    │ locking strength: for share
+    │
+    └── • project
+        │ columns: (x, y)
+        │
+        └── • index join
+            │ columns: (x, y, z)
+            │ estimated row count: 10 (missing stats)
+            │ table: xyz@xyz_pkey
+            │ key columns: x
+            │
+            └── • scan
+                  columns: (x, z)
+                  estimated row count: 10 (missing stats)
+                  table: xyz@xyz_z_idx
+                  spans: /2-/3
+
+# We can push down locking if the index join fetches all families.
 query T
 EXPLAIN (VERBOSE) SELECT * FROM xyz WHERE z = 2 FOR SHARE
 ----
 distribution: local
 vectorized: true
 ·
-• lookup join (semi)
+• index join
 │ columns: (x, y, z)
 │ estimated row count: 10 (missing stats)
 │ table: xyz@xyz_pkey
-│ equality: (x) = (x)
-│ equality cols are key
+│ key columns: x
 │ locking strength: for share
 │
-└── • index join
-    │ columns: (x, y, z)
-    │ estimated row count: 10 (missing stats)
-    │ table: xyz@xyz_pkey
-    │ key columns: x
-    │
-    └── • scan
-          columns: (x, z)
-          estimated row count: 10 (missing stats)
-          table: xyz@xyz_z_idx
-          spans: /2-/3
+└── • scan
+      columns: (x, z)
+      estimated row count: 10 (missing stats)
+      table: xyz@xyz_z_idx
+      spans: /2-/3
 
+# We can't push down locking if the lookup join doesn't fetch all families.
 query T
-EXPLAIN (VERBOSE) SELECT * FROM t INNER LOOKUP JOIN xyz ON x = b WHERE a = 5 FOR UPDATE OF xyz
+EXPLAIN (VERBOSE) SELECT t.*, x FROM t INNER LOOKUP JOIN xyz ON x = b WHERE a = 5 FOR UPDATE OF xyz
 ----
 distribution: local
 vectorized: true
 ·
 • lookup join (semi)
-│ columns: (a, b, x, y, z)
+│ columns: (a, b, x)
 │ estimated row count: 1 (missing stats)
 │ table: xyz@xyz_pkey
 │ equality: (x) = (x)
 │ equality cols are key
 │ locking strength: for update
 │
 └── • lookup join (inner)
-    │ columns: (a, b, x, y, z)
+    │ columns: (a, b, x)
     │ estimated row count: 1 (missing stats)
     │ table: xyz@xyz_pkey
     │ equality: (b) = (x)
@@ -3016,5 +3063,26 @@ vectorized: true
           table: t@t_pkey
           spans: /5/0
 
+# We can push down locking if the lookup join fetches all families.
+query T
+EXPLAIN (VERBOSE) SELECT * FROM t INNER LOOKUP JOIN xyz ON x = b WHERE a = 5 FOR UPDATE OF xyz
+----
+distribution: local
+vectorized: true
+·
+• lookup join (inner)
+│ columns: (a, b, x, y, z)
+│ estimated row count: 1 (missing stats)
+│ table: xyz@xyz_pkey
+│ equality: (b) = (x)
+│ equality cols are key
+│ locking strength: for update
+│
+└── • scan
+      columns: (a, b)
+      estimated row count: 1 (missing stats)
+      table: t@t_pkey
+      spans: /5/0
+
 statement ok
 RESET optimizer_use_lock_op_for_serializable

pkg/sql/opt/util.go

@@ -160,3 +160,17 @@ func GetAllFKsAmongTables(
 	}
 	return addFKs
 }
+
+// FamiliesForCols returns the set of column families for the set of cols.
+func FamiliesForCols(tab cat.Table, tabID TableID, cols ColSet) (families intsets.Fast) {
+	for i, n := 0, tab.FamilyCount(); i < n; i++ {
+		family := tab.Family(i)
+		for j, m := 0, family.ColumnCount(); j < m; j++ {
+			if cols.Contains(tabID.ColumnID(family.Column(j).Ordinal)) {
+				families.Add(i)
+				break
+			}
+		}
+	}
+	return families
+}