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 if the provided scan does not
output all column in the table scope. A scan on the table and its
logical properties are required in order to fully normalize partial
index predicates.

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

Release note: None

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

@@ -547,7 +547,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

@@ -622,7 +622,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
@@ -633,23 +633,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 {
-	_, 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)
 	pred, ok := PartialIndexPredicate(tabMeta, s.Index)
 	if !ok {
-		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 pred
 }

pkg/sql/opt/memo/expr_format.go

@@ -1252,7 +1252,7 @@ func FormatPrivate(f *ExprFmtCtx, private interface{}, physProps *physical.Requi
 		if ScanIsReverseFn(f.Memo.Metadata(), t, &physProps.Ordering) {
 			f.Buffer.WriteString(",rev")
 		}
-		if t.UsesPartialIndex(f.Memo.Metadata()) {
+		if t.PartialIndexPredicate(f.Memo.Metadata()) != nil {
 			f.Buffer.WriteString(",partial")
 		}
 

pkg/sql/opt/memo/logical_props_builder.go

@@ -62,12 +62,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
 	// ------------
@@ -91,7 +86,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)
@@ -117,7 +112,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
@@ -160,7 +155,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

@@ -622,11 +622,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

@@ -522,25 +522,8 @@ func (b *Builder) buildScan(
 
 		// 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)
-		}
+		// predicates.
+		b.addPartialIndexPredicatesForTable(tabMeta, outScope.expr)
 
 		if b.trackViewDeps {
 			dep := opt.ViewDep{DataSource: tab}
@@ -682,7 +665,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.
@@ -700,6 +688,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/join_order_builder_test.go

@@ -453,7 +453,7 @@ func parseVertexSet(sesStr string) vertexSet {
 func printVertexSet(set vertexSet) string {
 	buf := bytes.Buffer{}
 	for idx, ok := set.next(0); ok; idx, ok = set.next(idx + 1) {
-		buf.WriteString(string('A' + idx))
+		buf.WriteString(string(rune('A' + idx)))
 	}
 	return buf.String()
 }

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

@@ -140,7 +140,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.hasRejectFlag(rejectPartialIndexes) && isPartialIndex {
@@ -154,26 +154,19 @@ func (it *scanIndexIter) ForEachStartingAfter(ord int, f enumerateIndexFunc) {
 
 		filters := it.originalFilters
 
-		// If the index is a partial index, check whether or not the
-		// originalFilters imply the predicate.
+		// If the index is a partial index, check whether the filters imply the
+		// predicate.
 		if isPartialIndex {
-			pred, ok := memo.PartialIndexPredicate(it.tabMeta, 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
-			}
+			predFilters := *pred.(*memo.FiltersExpr)
 
-			// If there are no originalFilters, then skip over any partial
-			// indexes that are not pseudo-partial indexes.
-			if filters == nil && !pred.IsTrue() {
+			// If there are no filters, then skip over any partial indexes that
+			// are not pseudo-partial indexes.
+			if filters == nil && !predFilters.IsTrue() {
 				continue
 			}
 
 			if filters != nil {
-				remainingFilters, ok := it.im.FiltersImplyPredicate(filters, pred)
+				remainingFilters, ok := it.im.FiltersImplyPredicate(filters, predFilters)
 				if !ok {
 					// The originalFilters do not imply the predicate, so skip
 					// over the partial index.