opt: build partial index predicates only when all table columns are i…

…n scope

This commit fixes a panic induced by trying to build partial index
predicate expressions without all of a table's columns in-scope. Some
scans, like scans built for foreign key checks, do not include all of a
table's column in their scope. For such scans, optbuilder no longer
attempts to build a partial index predicate because the predicate may
reference a column that is not in scope.

As a result of this change, `opt.TableMeta` may not have a predicate
expression for all partial indexes. The `memo.PartialIndexPredicate`
function which retrieves the predicate expressions has been updated to
account for this case.

Fixes cockroachdb#55672

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

Release note (bug fix): An INSERT into a table with a foreign key
reference to a table with a partial index no longer causes an error.

pkg/sql/logictest/testdata/logic_test/partial_index

@@ -1298,3 +1298,31 @@ INSERT INTO t55387 VALUES (1, 1, 5);
 query I rowsort
 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.
+
+statement ok
+CREATE TABLE t55672_a (
+    a INT PRIMARY KEY,
+    t TIMESTAMPTZ DEFAULT NULL,
+    UNIQUE INDEX (a) WHERE t is NULL
+)
+
+statement ok
+CREATE TABLE t55672_b (
+    b INT PRIMARY KEY,
+    a INT NOT NULL REFERENCES t55672_a (a)
+)
+
+statement ok
+INSERT INTO t55672_a (a) VALUES (1)
+
+statement ok
+INSERT INTO t55672_b (b,a) VALUES (1,1)
+
+statement ok
+INSERT INTO t55672_a (a, t) VALUES (2, now())
+
+statement ok
+INSERT INTO t55672_b (b,a) VALUES (2,2)

pkg/sql/opt/memo/expr.go

@@ -642,11 +642,16 @@ func (s *ScanPrivate) UsesPartialIndex(md *opt.Metadata) bool {
 
 // PartialIndexPredicate returns the FiltersExpr representing the predicate of
 // the partial index that the scan uses. If the scan does not use a partial
-// index, this function panics. UsesPartialIndex should be called first to
-// determine if the scan operates over a partial index.
+// index or if a partial index predicate was not built for this index, this
+// function panics. UsesPartialIndex should be called first to determine if the
+// scan operates over a partial index.
 func (s *ScanPrivate) PartialIndexPredicate(md *opt.Metadata) FiltersExpr {
 	tabMeta := md.TableMeta(s.Table)
-	return PartialIndexPredicate(tabMeta, s.Index)
+	pred, ok := PartialIndexPredicate(tabMeta, s.Index)
+	if !ok {
+		panic(errors.AssertionFailedf("partial index predicate not found for %s", tabMeta.Table.Index(s.Index).Name()))
+	}
+	return pred
 }
 
 // UsesPartialIndex returns true if the the LookupJoinPrivate looks-up via a
@@ -906,14 +911,19 @@ func OutputColumnIsAlwaysNull(e RelExpr, col opt.ColumnID) bool {
 
 // PartialIndexPredicate returns the FiltersExpr representing the partial index
 // predicate at the given index ordinal. If the index at the ordinal is not a
-// partial index, this function panics. cat.Index.Predicate should be used first
-// to determine if the index is a partial index.
+// partial index or the predicate expression was not built for the index,
+// ok=false is returned.
 //
 // Note that TableMeta.PartialIndexPredicates is not initialized for mutation
-// queries. This function will panic in the context of a mutation.
-func PartialIndexPredicate(tabMeta *opt.TableMeta, ord cat.IndexOrdinal) FiltersExpr {
-	p := tabMeta.PartialIndexPredicates[ord]
-	return *p.(*FiltersExpr)
+// queries. This function will never return ok=true in the context of a mutation.
+func PartialIndexPredicate(
+	tabMeta *opt.TableMeta, ord cat.IndexOrdinal,
+) (pred FiltersExpr, ok bool) {
+	p, ok := tabMeta.PartialIndexPredicates[ord]
+	if !ok {
+		return nil, false
+	}
+	return *p.(*FiltersExpr), true
 }
 
 // FKCascades stores metadata necessary for building cascading queries.

pkg/sql/opt/optbuilder/mutation_builder_fk.go

@@ -135,7 +135,7 @@ func (mb *mutationBuilder) buildFKChecksAndCascadesForDelete() {
 		//  - with Cascade/SetNull/SetDefault, we create a cascading mutation to
 		//    modify or delete "orphaned" rows in the child table.
 		//  - with Restrict/NoAction, we create a check that causes an error if
-		//    there are any "orhpaned" rows in the child table.
+		//    there are any "orphaned" rows in the child table.
 		if a := h.fk.DeleteReferenceAction(); a != tree.Restrict && a != tree.NoAction {
 			telemetry.Inc(sqltelemetry.ForeignKeyCascadesUseCounter)
 			var builder memo.CascadeBuilder

pkg/sql/opt/optbuilder/select.go

@@ -521,8 +521,26 @@ func (b *Builder) buildScan(
 		outScope.expr = b.factory.ConstructScan(&private)
 
 		// Add the partial indexes after constructing the scan so we can use the
-		// logical properties of the scan to fully normalize the index predicates.
-		b.addPartialIndexPredicatesForTable(tabMeta, outScope)
+		// logical properties of the scan to fully normalize the index
+		// predicates. Partial index predicates are only added if the outScope
+		// contains all the table's ordinary columns. If it does not, partial
+		// index predicates cannot be built because they may reference columns
+		// not in outScope. In the most common case, the outScope has the same
+		// number of columns as the table and we can skip checking the each
+		// ordinary column exists in outScope.
+		containsAllOrdinaryTableColumns := true
+		if len(outScope.cols) != tab.ColumnCount() {
+			for i := 0; i < tab.ColumnCount(); i++ {
+				col := tab.Column(i)
+				if col.Kind() == cat.Ordinary && !outScope.colSet().Contains(tabID.ColumnID(col.Ordinal())) {
+					containsAllOrdinaryTableColumns = false
+					break
+				}
+			}
+		}
+		if containsAllOrdinaryTableColumns {
+			b.addPartialIndexPredicatesForTable(tabMeta, outScope)
+		}
 
 		if b.trackViewDeps {
 			dep := opt.ViewDep{DataSource: tab}

pkg/sql/opt/xform/custom_funcs.go

@@ -99,7 +99,13 @@ func (c *CustomFuncs) GenerateIndexScans(grp memo.RelExpr, scanPrivate *memo.Sca
 		// for an unconstrained index scan.
 		_, isPartialIndex := md.Table(scanPrivate.Table).Index(iter.IndexOrdinal()).Predicate()
 		if isPartialIndex {
-			pred := memo.PartialIndexPredicate(tabMeta, iter.IndexOrdinal())
+			pred, ok := memo.PartialIndexPredicate(tabMeta, iter.IndexOrdinal())
+			if !ok {
+				// A partial index predicate expression was not built for the
+				// partial index. Implication cannot be proven so it must be
+				// skipped.
+				continue
+			}
 			if !pred.IsTrue() {
 				continue
 			}
@@ -215,7 +221,13 @@ func (c *CustomFuncs) GeneratePartialIndexScans(
 	// Iterate over all partial indexes.
 	iter := makeScanIndexIter(c.e.mem, scanPrivate, rejectNonPartialIndexes)
 	for iter.Next() {
-		pred := memo.PartialIndexPredicate(tabMeta, iter.IndexOrdinal())
+		pred, ok := memo.PartialIndexPredicate(tabMeta, iter.IndexOrdinal())
+		if !ok {
+			// A partial index predicate expression was not built for the
+			// partial index. Implication cannot be proven so it must be
+			// skipped.
+			continue
+		}
 		remainingFilters, ok := c.im.FiltersImplyPredicate(filters, pred)
 		if !ok {
 			// The filters do not imply the predicate, so the partial index
@@ -346,9 +358,14 @@ func (c *CustomFuncs) GenerateConstrainedScans(
 		// If the index is a partial index, check whether or not the filter
 		// implies the predicate.
 		_, isPartialIndex := md.Table(scanPrivate.Table).Index(iter.IndexOrdinal()).Predicate()
-		var pred memo.FiltersExpr
 		if isPartialIndex {
-			pred = memo.PartialIndexPredicate(tabMeta, iter.IndexOrdinal())
+			pred, ok := memo.PartialIndexPredicate(tabMeta, iter.IndexOrdinal())
+			if !ok {
+				// A partial index predicate expression was not built for the
+				// partial index. Implication cannot be proven so it must be
+				// skipped.
+				continue
+			}
 			remainingFilters, ok := c.im.FiltersImplyPredicate(filters, pred)
 			if !ok {
 				// The filters do not imply the predicate, so the partial index
@@ -1727,7 +1744,13 @@ func (c *CustomFuncs) GenerateLookupJoins(
 		// ON filters.
 		_, isPartialIndex := md.Table(scanPrivate.Table).Index(iter.IndexOrdinal()).Predicate()
 		if isPartialIndex {
-			pred := memo.PartialIndexPredicate(tabMeta, iter.IndexOrdinal())
+			pred, ok := memo.PartialIndexPredicate(tabMeta, iter.IndexOrdinal())
+			if !ok {
+				// A partial index predicate expression was not built for the
+				// partial index. Implication cannot be proven so it must be
+				// skipped.
+				continue
+			}
 			remainingFilters, ok := c.im.FiltersImplyPredicate(onFilters, pred)
 			if !ok {
 				// The ON filters do not imply the predicate, so the partial

pkg/sql/opt/xform/testdata/rules/scan

@@ -748,3 +748,54 @@ scan with_time_index@secondary
  ├── constraint: /2/1: [/'2017-05-10 00:00:00' - ]
  ├── key: (1)
  └── fd: (1)-->(2)
+
+exec-ddl
+CREATE TABLE fk_a (
+    a INT PRIMARY KEY,
+    t TIMESTAMPTZ DEFAULT NULL,
+    UNIQUE INDEX (a) WHERE t is NULL
+)
+----
+
+exec-ddl
+CREATE TABLE fk_b (
+    b INT PRIMARY KEY,
+    a INT NOT NULL REFERENCES fk_a (a)
+)
+----
+
+# Do not use a non-implied partial index for FK check scans.
+opt expect-not=(GenerateIndexScans, GeneratePartialIndexScans)
+INSERT INTO fk_b (b,a) VALUES (1,1)
+----
+insert fk_b
+ ├── columns: <none>
+ ├── insert-mapping:
+ │    ├── column1:4 => b:1
+ │    └── column2:5 => fk_b.a:2
+ ├── input binding: &1
+ ├── cardinality: [0 - 0]
+ ├── volatile, mutations
+ ├── values
+ │    ├── columns: column1:4!null column2:5!null
+ │    ├── cardinality: [1 - 1]
+ │    ├── key: ()
+ │    ├── fd: ()-->(4,5)
+ │    └── (1, 1)
+ └── f-k-checks
+      └── f-k-checks-item: fk_b(a) -> fk_a(a)
+           └── anti-join (lookup fk_a)
+                ├── columns: column2:6!null
+                ├── key columns: [6] = [7]
+                ├── lookup columns are key
+                ├── cardinality: [0 - 1]
+                ├── key: ()
+                ├── fd: ()-->(6)
+                ├── with-scan &1
+                │    ├── columns: column2:6!null
+                │    ├── mapping:
+                │    │    └──  column2:5 => column2:6
+                │    ├── cardinality: [1 - 1]
+                │    ├── key: ()
+                │    └── fd: ()-->(6)
+                └── filters (true)