opt: build all partial index predicate expressions in TableMeta

Previously, partial index predicates were only built and added to
`TableMeta.PartialIndexPredicates` if the scope of the scan in
`Builder.buildScan` included all ordinary table columns. This behavior
prevented errors when building partial index predicate expressions for
foriegn key check scans. These scans do not include all columns, so when
building a predicate expression that referenced a table column not
included, optbuilder would err with "column does not exist".

This commit changes this behavior so that partial index predicates are
always built and added to `TableMeta.PartialIndexPredicates`. In order
to do this, `Builder.addPartialIndexPredicatesForTable` creates its own
table scope. It also constructs a new scan rather than relying on the
already-built scan. Constructing a scan is required for considering
logical properties while normalizing partial index predicates.

See discussion at cockroachdb#57622
for more context on this change.

Release note: None

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

@@ -565,7 +565,7 @@ func (b *Builder) buildScan(scan *memo.ScanExpr) (execPlan, error) {
 	md := b.mem.Metadata()
 	tab := md.Table(scan.Table)
 
-	if !b.disableTelemetry && scan.UsesPartialIndex(md) {
+	if !b.disableTelemetry && scan.PartialIndexPredicate(md) != nil {
 		telemetry.Inc(sqltelemetry.PartialIndexScanUseCounter)
 	}
 

pkg/sql/opt/memo/expr.go

@@ -639,7 +639,7 @@ func (s *ScanPrivate) IsUnfiltered(md *opt.Metadata) bool {
 	return (s.Constraint == nil || s.Constraint.IsUnconstrained()) &&
 		s.InvertedConstraint == nil &&
 		s.HardLimit == 0 &&
-		!s.UsesPartialIndex(md)
+		s.PartialIndexPredicate(md) == nil
 }
 
 // IsLocking returns true if the ScanPrivate is configured to use a row-level
@@ -650,29 +650,15 @@ func (s *ScanPrivate) IsLocking() bool {
 	return s.Locking != nil
 }
 
-// UsesPartialIndex returns true if the ScanPrivate indicates a scan over a
-// partial index.
-func (s *ScanPrivate) UsesPartialIndex(md *opt.Metadata) bool {
-	if s.Index == cat.PrimaryIndex {
-		// Primary index is always non-partial; skip making the catalog calls.
-		return false
-	}
-	_, isPartialIndex := md.Table(s.Table).Index(s.Index).Predicate()
-	return isPartialIndex
-}
-
 // PartialIndexPredicate returns the FiltersExpr representing the predicate of
 // the partial index that the scan uses. If the scan does not use a partial
-// 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.
+// index, nil is returned.
 func (s *ScanPrivate) PartialIndexPredicate(md *opt.Metadata) FiltersExpr {
 	tabMeta := md.TableMeta(s.Table)
 	p, ok := tabMeta.PartialIndexPredicates[s.Index]
 	if !ok {
-		// A partial index predicate expression was not built for the
-		// partial index.
-		panic(errors.AssertionFailedf("partial index predicate not found for %s", tabMeta.Table.Index(s.Index).Name()))
+		// The index is not a partial index.
+		return nil
 	}
 	return *p.(*FiltersExpr)
 }

pkg/sql/opt/memo/logical_props_builder.go

@@ -63,12 +63,7 @@ func (b *logicalPropsBuilder) clear() {
 func (b *logicalPropsBuilder) buildScanProps(scan *ScanExpr, rel *props.Relational) {
 	md := scan.Memo().Metadata()
 	hardLimit := scan.HardLimit.RowCount()
-
-	isPartialIndexScan := scan.UsesPartialIndex(md)
-	var pred FiltersExpr
-	if isPartialIndexScan {
-		pred = scan.PartialIndexPredicate(md)
-	}
+	pred := scan.PartialIndexPredicate(md)
 
 	// Side Effects
 	// ------------
@@ -92,7 +87,7 @@ func (b *logicalPropsBuilder) buildScanProps(scan *ScanExpr, rel *props.Relation
 	}
 	// Union not-NULL columns with not-NULL columns in the partial index
 	// predicate.
-	if isPartialIndexScan {
+	if pred != nil {
 		rel.NotNullCols.UnionWith(b.rejectNullCols(pred))
 	}
 	rel.NotNullCols.IntersectionWith(rel.OutputCols)
@@ -118,7 +113,7 @@ func (b *logicalPropsBuilder) buildScanProps(scan *ScanExpr, rel *props.Relation
 		if tabMeta := md.TableMeta(scan.Table); tabMeta.Constraints != nil {
 			b.addFiltersToFuncDep(*tabMeta.Constraints.(*FiltersExpr), &rel.FuncDeps)
 		}
-		if isPartialIndexScan {
+		if pred != nil {
 			b.addFiltersToFuncDep(pred, &rel.FuncDeps)
 
 			// Partial index keys are not added to the functional dependencies in
@@ -161,7 +156,7 @@ func (b *logicalPropsBuilder) buildScanProps(scan *ScanExpr, rel *props.Relation
 		if scan.Constraint != nil {
 			b.updateCardinalityFromConstraint(scan.Constraint, rel)
 		}
-		if isPartialIndexScan {
+		if pred != nil {
 			b.updateCardinalityFromFilters(pred, rel)
 		}
 	}

pkg/sql/opt/memo/statistics_builder.go

@@ -623,11 +623,7 @@ func (sb *statisticsBuilder) buildScan(scan *ScanExpr, relProps *props.Relationa
 
 	inputStats := sb.makeTableStatistics(scan.Table)
 	s.RowCount = inputStats.RowCount
-
-	var pred FiltersExpr
-	if scan.UsesPartialIndex(sb.md) {
-		pred = scan.PartialIndexPredicate(sb.md)
-	}
+	pred := scan.PartialIndexPredicate(sb.md)
 
 	// If the constraints and pred are nil, then this scan is an unconstrained
 	// scan on a non-partial index. The stats of the scan are the same as the

pkg/sql/opt/optbuilder/partial_index.go

@@ -23,14 +23,14 @@ import (
 // Returns an error if any non-immutable operators are found.
 //
 // Note: This function should only be used to build partial index or arbiter
-// predicate expressions that have only a table's columns in scope and that are
-// not part of the relational expression tree. For example, this is used to
-// populate the TableMeta.PartialIndexPredicates cache and for determining
-// arbiter indexes in UPSERT and INSERT ON CONFLICT mutations. But it is not
-// used for building synthesized mutation columns that determine whether or not
-// to PUT or DEL a partial index entry for a row; these synthesized columns are
-// projected as part of the opt expression tree and they reference columns
-// beyond a table's base scope.
+// predicate expressions that have only a table's ordinary columns in scope and
+// that are not part of the relational expression tree. For example, this is
+// used to populate the TableMeta.PartialIndexPredicates cache and for
+// determining arbiter indexes in UPSERT and INSERT ON CONFLICT mutations. But
+// it is not used for building synthesized mutation columns that determine
+// whether to issue PUT or DEL operations on a partial index for a mutated row;
+// these synthesized columns are projected as part of the opt expression tree
+// and they can reference columns not part of a table's ordinary columns.
 func (b *Builder) buildPartialIndexPredicate(
 	tableScope *scope, expr tree.Expr, context string,
 ) (memo.FiltersExpr, error) {

pkg/sql/opt/optbuilder/select.go

@@ -532,6 +532,7 @@ func (b *Builder) buildScan(
 
 	b.addCheckConstraintsForTable(tabMeta)
 	b.addComputedColsForTable(tabMeta)
+	b.addPartialIndexPredicatesForTable(tabMeta)
 
 	outScope.expr = b.factory.ConstructScan(&private)
 
@@ -551,28 +552,6 @@ func (b *Builder) buildScan(
 		outScope.expr = b.factory.ConstructProject(outScope.expr, proj, tabColIDs)
 	}
 
-	// Add the partial indexes after constructing the scan so we can use the
-	// 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 that 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}
 		dep.ColumnIDToOrd = make(map[opt.ColumnID]int)
@@ -712,7 +691,7 @@ func (b *Builder) addComputedColsForTable(tabMeta *opt.TableMeta) {
 //
 // The predicates are used as "known truths" about table data. Any predicates
 // containing non-immutable operators are omitted.
-func (b *Builder) addPartialIndexPredicatesForTable(tabMeta *opt.TableMeta, tableScope *scope) {
+func (b *Builder) addPartialIndexPredicatesForTable(tabMeta *opt.TableMeta) {
 	tab := tabMeta.Table
 
 	// Find the first partial index.
@@ -730,6 +709,15 @@ func (b *Builder) addPartialIndexPredicatesForTable(tabMeta *opt.TableMeta, tabl
 		return
 	}
 
+	// Construct a scan as the tableScope expr so that logical properties of the
+	// scan can be used to fully normalize the index predicate.
+	tableScope := b.allocScope()
+	tableScope.appendOrdinaryColumnsFromTable(tabMeta, &tabMeta.Alias)
+	tableScope.expr = b.factory.ConstructScan(&memo.ScanPrivate{
+		Table: tabMeta.MetaID,
+		Cols:  tableScope.colSet(),
+	})
+
 	// Skip to the first partial index we found above.
 	for ; indexOrd < numIndexes; indexOrd++ {
 		index := tab.Index(indexOrd)

pkg/sql/opt/xform/limit_funcs.go

@@ -54,7 +54,8 @@ func (c *CustomFuncs) CanLimitFilteredScan(
 		return false
 	}
 
-	if scanPrivate.Constraint == nil && !scanPrivate.UsesPartialIndex(c.e.mem.Metadata()) {
+	md := c.e.mem.Metadata()
+	if scanPrivate.Constraint == nil && scanPrivate.PartialIndexPredicate(md) == nil {
 		// This is not a constrained scan nor a partial index scan, so skip it.
 		// The GenerateLimitedScans rule is responsible for limited
 		// unconstrained scans on non-partial indexes.

pkg/sql/opt/xform/scan_index_iter.go

@@ -198,7 +198,7 @@ func (it *scanIndexIter) ForEachStartingAfter(ord int, f enumerateIndexFunc) {
 			continue
 		}
 
-		_, isPartialIndex := index.Predicate()
+		pred, isPartialIndex := it.tabMeta.PartialIndexPredicates[ord]
 
 		// Skip over partial indexes if rejectPartialIndexes is set.
 		if it.hasRejectFlags(rejectPartialIndexes) && isPartialIndex {
@@ -212,27 +212,19 @@ func (it *scanIndexIter) ForEachStartingAfter(ord int, f enumerateIndexFunc) {
 
 		filters := it.filters
 
-		// If the index is a partial index, check whether or not the filters
-		// imply the predicate.
+		// If the index is a partial index, check whether the filters imply the
+		// predicate.
 		if isPartialIndex {
-			p, ok := it.tabMeta.PartialIndexPredicates[ord]
-			if !ok {
-				// A partial index predicate expression was not built for the
-				// partial index. See Builder.buildScan for details on when this
-				// can occur. Implication cannot be proven so it must be
-				// skipped.
-				continue
-			}
-			pred := *p.(*memo.FiltersExpr)
+			p := *pred.(*memo.FiltersExpr)
 
 			// If there are no filters, then skip over any partial indexes that
 			// are not pseudo-partial indexes.
-			if filters == nil && !pred.IsTrue() {
+			if filters == nil && !p.IsTrue() {
 				continue
 			}
 
 			if filters != nil {
-				remainingFilters, ok := it.filtersImplyPredicate(pred)
+				remainingFilters, ok := it.filtersImplyPredicate(p)
 				if !ok {
 					// The predicate is not implied by the filters, so skip over
 					// the partial index.

pkg/sql/opt/xform/testdata/rules/select

@@ -246,7 +246,7 @@ CREATE INDEX idx2 ON p (s) WHERE i > 0
 memo expect=GeneratePartialIndexScans
 SELECT * FROM p WHERE i > 0 AND s = 'foo'
 ----
-memo (optimized, ~13KB, required=[presentation: i:1,f:2,s:3,b:4])
+memo (optimized, ~15KB, required=[presentation: i:1,f:2,s:3,b:4])
  ├── G1: (select G2 G3) (index-join G4 p,cols=(1-4)) (index-join G5 p,cols=(1-4)) (index-join G6 p,cols=(1-4)) (index-join G7 p,cols=(1-4))
  │    └── [presentation: i:1,f:2,s:3,b:4]
  │         ├── best: (index-join G4 p,cols=(1-4))
@@ -294,7 +294,7 @@ memo (optimized, ~13KB, required=[presentation: i:1,f:2,s:3,b:4])
 memo expect-not=GeneratePartialIndexScans
 SELECT i FROM p WHERE s = 'bar'
 ----
-memo (optimized, ~8KB, required=[presentation: i:1])
+memo (optimized, ~9KB, required=[presentation: i:1])
  ├── G1: (project G2 G3 i)
  │    └── [presentation: i:1]
  │         ├── best: (project G2 G3 i)
@@ -1493,7 +1493,7 @@ CREATE INDEX idx2 ON p (i) WHERE s = 'foo'
 memo
 SELECT i FROM p WHERE i = 3 AND s = 'foo'
 ----
-memo (optimized, ~17KB, required=[presentation: i:1])
+memo (optimized, ~19KB, required=[presentation: i:1])
  ├── G1: (project G2 G3 i) (project G4 G3 i) (project G5 G3 i)
  │    └── [presentation: i:1]
  │         ├── best: (project G5 G3 i)
@@ -1566,7 +1566,7 @@ CREATE INDEX idx ON p (i) WHERE s = 'foo'
 memo expect-not=GenerateConstrainedScans
 SELECT i FROM p WHERE s = 'bar' AND i = 5
 ----
-memo (optimized, ~7KB, required=[presentation: i:1])
+memo (optimized, ~8KB, required=[presentation: i:1])
  ├── G1: (project G2 G3 i)
  │    └── [presentation: i:1]
  │         ├── best: (project G2 G3 i)