opt: fix FK cascading updates to child tables with partial indexes

This commit fixes several bugs that are present when using foreign key
cascading updates and partial indexes.

Previously, the optimizer was not synthesizing partial index DEL columns
for FK cascading updates. As a result, a partial index on a child table
could become inconsistent with the rows in the primary index,
ultimately resulting in incorrect query results. The optbuilder has
been refactored to project these columns and to reduce the complexity
of doing so. As a result, partial index PUT and DEL columns are now
projected in the same expression, rather than the DEL columns being
projected as far down in the expression tree as possible.

Additionally, the execution engine was unable to handle extraneous
columns that can be added as input to FK cascading updates. These
extraneous columns would be incorrectly interpreted as synthesized
partial index columns. This commit works around this issue by passing a
set of partial indexes that columns have been synthesized for from the
optimizer to the execution engine. The longer term fix is to not produce
these columns (see issue cockroachdb#57097).

Fixes cockroachdb#57085
Fixes cockroachdb#57084

Release justification: This is a critical bug fix to a new feature,
partial indexes.

Release note (bug fix): Previously, updating parent table of a foreign
key relationship with cascading updates could cause errors or
inconsistencies in partial indexes for child tables with partial
indexes. The inconsistency of partial indexes could result in incorrect
query results. This has been fixed.

pkg/sql/distsql_spec_exec_factory.go

@@ -819,6 +819,7 @@ func (e *distSQLSpecExecFactory) ConstructUpdate(
 	updateCols exec.TableColumnOrdinalSet,
 	returnCols exec.TableColumnOrdinalSet,
 	checks exec.CheckOrdinalSet,
+	partialIndexes exec.PartialIndexOrdinalSet,
 	passthrough colinfo.ResultColumns,
 	autoCommit bool,
 ) (exec.Node, error) {

pkg/sql/logictest/testdata/logic_test/partial_index

@@ -1426,6 +1426,7 @@ SELECT k FROM t55387 WHERE a > 1 AND b > 3
 
 # Regression test for #55672. Do not build partial index predicates when the
 # scope does not include all table columns, like FK check scans.
+subtest regression_55672
 
 statement ok
 CREATE TABLE t55672_a (
@@ -1451,3 +1452,84 @@ INSERT INTO t55672_a (a, t) VALUES (2, now())
 
 statement ok
 INSERT INTO t55672_b (b,a) VALUES (2,2)
+
+# Regression test for #57085. Cascading UPDATEs should correctly update partial
+# indexes of the child table.
+subtest regression_57085
+
+# Update a partial index in a child table.
+statement ok
+CREATE TABLE t57085a_groups (
+    id INT PRIMARY KEY
+);
+CREATE TABLE t57085a_users (
+    id INT PRIMARY KEY,
+    group_id INT REFERENCES t57085a_groups ON UPDATE CASCADE,
+    i INT,
+    INDEX idx (group_id) WHERE i > 0
+);
+
+statement ok
+INSERT INTO t57085a_groups VALUES (1);
+INSERT INTO t57085a_users VALUES (10, 1, 100), (20, 1, -100);
+UPDATE t57085a_groups SET id = 2 WHERE id = 1;
+
+query II rowsort
+SELECT id, i FROM t57085a_users@idx WHERE i > 0 AND group_id = 2
+----
+10  100
+
+# Update a partial index in a child table with a single variable boolean
+# predicate.
+statement ok
+CREATE TABLE t57085b_groups (
+    id INT PRIMARY KEY
+);
+CREATE TABLE t57085b_users (
+    id INT PRIMARY KEY,
+    group_id INT REFERENCES t57085b_groups ON UPDATE CASCADE,
+    b BOOL,
+    INDEX idx (group_id) WHERE b
+);
+
+statement ok
+INSERT INTO t57085b_groups VALUES (1);
+INSERT INTO t57085b_users VALUES (10, 1, true), (20, 1, false);
+UPDATE t57085b_groups SET id = 2 WHERE id = 1;
+
+query IB rowsort
+SELECT id, b FROM t57085b_users@idx WHERE b AND group_id = 2
+----
+10  true
+
+# Update the group's id with an INSERT ON CONFLICT DO UPDATE.
+statement ok
+INSERT INTO t57085b_groups VALUES (2) ON CONFLICT (id) DO UPDATE SET id = 3
+
+query IB rowsort
+SELECT id, b FROM t57085b_users@idx WHERE b AND group_id = 3
+----
+10  true
+
+# Update a partial index that references the column being updated in the
+# cascade.
+statement ok
+CREATE TABLE t57085c_groups (
+    id INT PRIMARY KEY
+);
+CREATE TABLE t57085c_users (
+    id INT PRIMARY KEY,
+    group_id INT REFERENCES t57085c_groups ON UPDATE CASCADE,
+    i INT,
+    INDEX idx (i) WHERE group_id = 3
+);
+
+statement ok
+INSERT INTO t57085c_groups VALUES (1), (2);
+INSERT INTO t57085c_users VALUES (10, 1, 100), (20, 2, 200);
+UPDATE t57085c_groups SET id = 3 WHERE id = 1;
+
+query II rowsort
+SELECT id, i FROM t57085c_users@idx WHERE group_id = 3 AND i = 100
+----
+10  100

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

@@ -44,6 +44,8 @@ func (b *Builder) buildMutationInput(
 		return execPlan{}, err
 	}
 
+	// TODO(mgartner/radu): This can incorrectly append columns in a FK cascade
+	// update that are never used during execution. See issue #57097.
 	if p.WithID != 0 {
 		// The input might have extra columns that are used only by FK or unique
 		// checks; make sure we don't project them away.
@@ -310,6 +312,17 @@ func (b *Builder) buildUpdate(upd *memo.UpdateExpr) (execPlan, error) {
 	updateColOrds := ordinalSetFromColList(upd.UpdateCols)
 	returnColOrds := ordinalSetFromColList(upd.ReturnCols)
 	checkOrds := ordinalSetFromColList(upd.CheckCols)
+	partialIndexPutOrds := ordinalSetFromColList(upd.PartialIndexPutCols)
+	partialIndexDelOrds := ordinalSetFromColList(upd.PartialIndexDelCols)
+
+	if partialIndexPutOrds.Len() != partialIndexDelOrds.Len() {
+		err := errors.AssertionFailedf(
+			"expected the same number of partial index PUT and DEL cols, got %d PUT columns and %d DEL columns",
+			partialIndexPutOrds.Len(),
+			partialIndexDelOrds.Len(),
+		)
+		return execPlan{}, err
+	}
 
 	// Construct the result columns for the passthrough set.
 	var passthroughCols colinfo.ResultColumns
@@ -327,6 +340,7 @@ func (b *Builder) buildUpdate(upd *memo.UpdateExpr) (execPlan, error) {
 		updateColOrds,
 		returnColOrds,
 		checkOrds,
+		partialIndexPutOrds,
 		passthroughCols,
 		b.allowAutoCommit && len(upd.UniqueChecks) == 0 &&
 			len(upd.FKChecks) == 0 && len(upd.FKCascades) == 0,

pkg/sql/opt/exec/factory.go

@@ -125,6 +125,10 @@ type NodeColumnOrdinalSet = util.FastIntSet
 // taken from the opt.Table.Check collection.
 type CheckOrdinalSet = util.FastIntSet
 
+// PartialIndexSet contains the ordinal positions of a set of partial indexes
+// from a table's collection of indexes.
+type PartialIndexOrdinalSet = util.FastIntSet
+
 // AggInfo represents an aggregation (see ConstructGroupBy).
 type AggInfo struct {
 	FuncName   string

pkg/sql/opt/exec/factory.opt

@@ -421,6 +421,7 @@ define Update {
     UpdateCols exec.TableColumnOrdinalSet
     ReturnCols exec.TableColumnOrdinalSet
     Checks exec.CheckOrdinalSet
+    PartialIndexes exec.PartialIndexOrdinalSet
     Passthrough colinfo.ResultColumns
 
     # If set, the operator will commit the transaction as part of its execution.

pkg/sql/opt/optbuilder/delete.go

@@ -76,6 +76,9 @@ func (b *Builder) buildDelete(del *tree.Delete, inScope *scope) (outScope *scope
 func (mb *mutationBuilder) buildDelete(returning tree.ReturningExprs) {
 	mb.buildFKChecksAndCascadesForDelete()
 
+	// Project partial index DEL boolean columns.
+	mb.projectPartialIndexCols(nil /* putScope */, mb.fetchScope)
+
 	private := mb.makeMutationPrivate(returning != nil)
 	mb.outScope.expr = mb.b.factory.ConstructDelete(
 		mb.outScope.expr, mb.uniqueChecks, mb.fkChecks, private,

pkg/sql/opt/optbuilder/fk_cascade.go

@@ -634,6 +634,9 @@ func (cb *onUpdateCascadeBuilder) Build(
 			cb.childTable, &mb.alias, fk, binding, bindingProps, oldValues, newValues,
 		)
 
+		// Set the fetchScope to the scope of the fetch expression.
+		mb.fetchScope = mb.outScope
+
 		// The scope created by b.buildUpdateCascadeMutationInput has the table
 		// columns, followed by the old FK values, followed by the new FK values.
 		numFKCols := fk.ColumnCount()

pkg/sql/opt/optbuilder/insert.go

@@ -649,8 +649,8 @@ func (mb *mutationBuilder) buildInsert(returning tree.ReturningExprs) {
 	// Add any check constraint boolean columns to the input.
 	mb.addCheckConstraintCols()
 
-	// Add any partial index put boolean columns to the input.
-	mb.projectPartialIndexPutCols(preCheckScope)
+	// Project partial index PUT boolean columns.
+	mb.projectPartialIndexCols(preCheckScope, nil /* delScope */)
 
 	mb.buildUniqueChecksForInsert()
 
@@ -918,6 +918,9 @@ func (mb *mutationBuilder) buildInputForUpsert(
 		inScope,
 	)
 
+	// Set the fetchScope to the scope of the fetch expression.
+	mb.fetchScope = fetchScope
+
 	// If the index is a unique partial index, then rows that are not in the
 	// partial index cannot conflict with insert rows. Therefore, a Select wraps
 	// the scan on the right side of the left outer join with the partial index
@@ -998,9 +1001,6 @@ func (mb *mutationBuilder) buildInputForUpsert(
 
 	mb.targetColList = make(opt.ColList, 0, mb.tab.ColumnCount())
 	mb.targetColSet = opt.ColSet{}
-
-	// Add any partial index del boolean columns to the input for UPSERTs.
-	mb.projectPartialIndexDelCols(fetchScope)
 }
 
 // setUpsertCols sets the list of columns to be updated in case of conflict.
@@ -1070,25 +1070,8 @@ func (mb *mutationBuilder) buildUpsert(returning tree.ReturningExprs) {
 	// Add any check constraint boolean columns to the input.
 	mb.addCheckConstraintCols()
 
-	// Add any partial index put boolean columns. The variables in these partial
-	// index predicates must resolve to the new column values of the row which
-	// are either the existing values of the columns or new values provided in
-	// the upsert. Therefore, the variables must resolve to the upsert CASE
-	// expression columns, so the project must be added after the upsert columns
-	// are.
-	//
-	// For example, consider the table and upsert:
-	//
-	//   CREATE TABLE t (a INT PRIMARY KEY, b INT, INDEX (b) WHERE a > 1)
-	//   INSERT INTO t (a, b) VALUES (1, 2) ON CONFLICT (a) DO UPDATE a = t.a + 1
-	//
-	// An entry in the partial index should only be added when a > 1. The
-	// resulting value of a is dependent on whether or not there is a conflict.
-	// In the case of no conflict, the (1, 2) is inserted into the table, and no
-	// partial index entry should be added. But if there is a conflict, The
-	// existing row where a = 1 has a incremented to 2, and an entry should be
-	// added to the partial index.
-	mb.projectPartialIndexPutCols(preCheckScope)
+	// Project partial index PUT and DEL boolean columns.
+	mb.projectPartialIndexCols(preCheckScope, mb.fetchScope)
 
 	mb.buildFKChecksForUpsert()
 

pkg/sql/opt/optbuilder/mutation_builder.go

@@ -56,6 +56,9 @@ type mutationBuilder struct {
 	// expression is completed, it will be contained in outScope.expr.
 	outScope *scope
 
+	// fetchScope contains the set of columns fetched from the target table.
+	fetchScope *scope
+
 	// targetColList is an ordered list of IDs of the table columns into which
 	// values will be inserted, or which will be updated with new values. It is
 	// incrementally built as the mutation operator is built.
@@ -335,8 +338,8 @@ func (mb *mutationBuilder) buildInputForUpdate(
 		}
 	}
 
-	// Add partial index del boolean columns to the input.
-	mb.projectPartialIndexDelCols(scanScope)
+	// Set the fetchScope to the scope of the fetch expression.
+	mb.fetchScope = mb.outScope
 }
 
 // buildInputForDelete constructs a Select expression from the fields in
@@ -402,8 +405,8 @@ func (mb *mutationBuilder) buildInputForDelete(
 	// Set list of columns that will be fetched by the input expression.
 	mb.setFetchColIDs(mb.outScope.cols)
 
-	// Add partial index boolean columns to the input.
-	mb.projectPartialIndexDelCols(scanScope)
+	// Set the fetchScope to the scope of the fetch expression.
+	mb.fetchScope = mb.outScope
 }
 
 // addTargetColsByName adds one target column for each of the names in the given
@@ -796,51 +799,45 @@ func (mb *mutationBuilder) mutationColumnIDs() opt.ColSet {
 	return cols
 }
 
-// projectPartialIndexPutCols builds a Project that synthesizes boolean output
-// columns for each partial index defined on the target table. The execution
-// code uses these booleans to determine whether or not to add a row in the
-// partial index.
-//
-// predScope is the scope of columns available to the partial index predicate
-// expression.
-func (mb *mutationBuilder) projectPartialIndexPutCols(predScope *scope) {
-	mb.projectPartialIndexCols(mb.partialIndexPutColIDs, predScope, "partial_index_put")
-}
-
-// projectPartialIndexPutCols builds a Project that synthesizes boolean output
-// columns for each partial index defined on the target table. The execution
-// code uses these booleans to determine whether or not to remove a row in the
-// partial index.
-//
-// predScope is the scope of columns available to the partial index predicate
-// expression.
-func (mb *mutationBuilder) projectPartialIndexDelCols(predScope *scope) {
-	mb.projectPartialIndexCols(mb.partialIndexDelColIDs, predScope, "partial_index_del")
-}
-
 // projectPartialIndexCols builds a Project that synthesizes boolean output
-// columns for each partial index defined on the target table.
-func (mb *mutationBuilder) projectPartialIndexCols(
-	colIDs opt.ColList, predScope *scope, aliasPrefix string,
-) {
+// columns for each partial index defined on the target table. PUT columns are
+// only projected if putScope is not nil and DEL columns are only projected if
+// delScope is not nil.
+func (mb *mutationBuilder) projectPartialIndexCols(putScope *scope, delScope *scope) {
 	if partialIndexCount(mb.tab) > 0 {
 		projectionScope := mb.outScope.replace()
 		projectionScope.appendColumnsFromScope(mb.outScope)
 
 		ord := 0
 		for i, n := 0, mb.tab.DeletableIndexCount(); i < n; i++ {
 			index := mb.tab.Index(i)
+
+			// Skip non-partial indexes.
 			if _, isPartial := index.Predicate(); !isPartial {
 				continue
 			}
 
 			expr := mb.parsePartialIndexPredicateExpr(i)
-			texpr := predScope.resolveAndRequireType(expr, types.Bool)
-			alias := fmt.Sprintf("%s%d", aliasPrefix, ord+1)
-			scopeCol := mb.b.addColumn(projectionScope, alias, texpr)
 
-			mb.b.buildScalar(texpr, predScope, projectionScope, scopeCol, nil)
-			colIDs[ord] = scopeCol.id
+			// Build a synthesized PUT columns.
+			if putScope != nil {
+				texpr := putScope.resolveAndRequireType(expr, types.Bool)
+				alias := fmt.Sprintf("partial_index_put%d", ord+1)
+				scopeCol := mb.b.addColumn(projectionScope, alias, texpr)
+
+				mb.b.buildScalar(texpr, putScope, projectionScope, scopeCol, nil)
+				mb.partialIndexPutColIDs[ord] = scopeCol.id
+			}
+
+			// Build a synthesized DEL columns.
+			if delScope != nil {
+				texpr := delScope.resolveAndRequireType(expr, types.Bool)
+				alias := fmt.Sprintf("partial_index_del%d", ord+1)
+				scopeCol := mb.b.addColumn(projectionScope, alias, texpr)
+
+				mb.b.buildScalar(texpr, delScope, projectionScope, scopeCol, nil)
+				mb.partialIndexDelColIDs[ord] = scopeCol.id
+			}
 
 			ord++
 		}