Merge #57707

57707: opt: build all partial index predicate expressions in TableMeta r=mgartner a=mgartner

#### 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 #57622
for more context on this change.

Release note: None

#### optbuilder: reduce redundant building of arbiter filter expressions

Previously, the arbiter predicate provided in the `WHERE` clause of an
`INSERT ... ON CONFLICT` statement was rebuilt as a `memo.FiltersExpr`
repetitively while the optbuilder determined arbiter indexes. Now, the
expression is built only once, reducing work.

Release note: None

Co-authored-by: Marcus Gartner <[email protected]>

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/insert.go

@@ -1213,6 +1213,7 @@ func (mb *mutationBuilder) arbiterIndexes(
 	tabMeta := mb.md.TableMeta(mb.tabID)
 	var tableScope *scope
 	var im *partialidx.Implicator
+	var arbiterFilters memo.FiltersExpr
 	for idx, idxCount := 0, mb.tab.IndexCount(); idx < idxCount; idx++ {
 		index := mb.tab.Index(idx)
 
@@ -1280,16 +1281,20 @@ func (mb *mutationBuilder) arbiterIndexes(
 				im.Init(mb.b.factory, mb.md, mb.b.evalCtx)
 			}
 
-			arbiterFilter, err := mb.b.buildPartialIndexPredicate(
-				tableScope, arbiterPredicate, "arbiter predicate",
-			)
-			if err != nil {
-				// The error is due to a non-immutable operator in the arbiter
-				// predicate. Continue on to see if a matching non-partial or
-				// pseudo-partial index exists.
-				continue
+			// Build the arbiter filters once.
+			if arbiterFilters == nil {
+				arbiterFilters, err = mb.b.buildPartialIndexPredicate(
+					tableScope, arbiterPredicate, "arbiter predicate",
+				)
+				if err != nil {
+					// The error is due to a non-immutable operator in the arbiter
+					// predicate. Continue on to see if a matching non-partial or
+					// pseudo-partial index exists.
+					continue
+				}
 			}
-			if _, ok := im.FiltersImplyPredicate(arbiterFilter, predFilter); ok {
+
+			if _, ok := im.FiltersImplyPredicate(arbiterFilters, predFilter); ok {
 				arbiters.Add(idx)
 			}
 		}

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

@@ -535,6 +535,11 @@ 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.expr)
+
 	if !virtualColIDs.Empty() {
 		// Project the expressions for the virtual columns (and pass through all
 		// scanned columns).
@@ -551,28 +556,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 +695,12 @@ 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) {
+//
+// scan is an optional argument that is a Scan expression on the table. If scan
+// outputs all the ordinary columns in the table, we avoid constructing a new
+// scan. A scan and its logical properties are required in order to fully
+// normalize the partial index predicates.
+func (b *Builder) addPartialIndexPredicatesForTable(tabMeta *opt.TableMeta, scan memo.RelExpr) {
 	tab := tabMeta.Table
 
 	// Find the first partial index.
@@ -730,6 +718,25 @@ 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)
+
+	// If the optional scan argument was provided and it outputs all of the
+	// ordinary table columns, we use it as tableScope.expr. Otherwise, we must
+	// construct a new scan. Attaching a scan to tableScope.expr is required to
+	// fully normalize the partial index predicates with logical properties of
+	// the scan.
+	if scan != nil && tableScope.colSet().SubsetOf(scan.Relational().OutputCols) {
+		tableScope.expr = scan
+	} else {
+		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)
+			predFilters := *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 && !predFilters.IsTrue() {
 				continue
 			}
 
 			if filters != nil {
-				remainingFilters, ok := it.filtersImplyPredicate(pred)
+				remainingFilters, ok := it.filtersImplyPredicate(predFilters)
 				if !ok {
 					// The predicate is not implied by the filters, so skip over
 					// the partial index.