opt: add helper functions for partial indexes and scans

This commit adds helper functions to determine if indexes are partial
indexes, if scans operate on partial indexes, and for retrieving partial
index predicate expressions.

Within the optimizer `TableMeta.PartialIndexPredicates` is now the
source of truth that should be used for answering these types of
questions. The catalog `Index.Predicate` function should only be used by
the optbuilder.

This commit also makes certain that there is an entry in the
`TableMeta.PartialIndexPredicates` map for every partial index on the
table, even if the partial index predicate is non-immutable. This makes
the map a safe source of truth for determining which indexes are
partial.

Release note: None

pkg/sql/opt/memo/expr.go

@@ -555,11 +555,10 @@ func (s *ScanPrivate) IsCanonical() bool {
 // IsUnfiltered returns true if the ScanPrivate will produce all rows in the
 // table.
 func (s *ScanPrivate) IsUnfiltered(md *opt.Metadata) bool {
-	_, isPartialIndex := md.Table(s.Table).Index(s.Index).Predicate()
-	return !isPartialIndex &&
-		(s.Constraint == nil || s.Constraint.IsUnconstrained()) &&
+	return (s.Constraint == nil || s.Constraint.IsUnconstrained()) &&
 		s.InvertedConstraint == nil &&
-		s.HardLimit == 0
+		s.HardLimit == 0 &&
+		!s.UsesPartialIndex(md)
 }
 
 // IsLocking returns true if the ScanPrivate is configured to use a row-level
@@ -570,6 +569,13 @@ 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 {
+	tabMeta := md.TableMeta(s.Table)
+	return IsPartialIndex(tabMeta, s.Index)
+}
+
 // NeedResults returns true if the mutation operator can return the rows that
 // were mutated.
 func (m *MutationPrivate) NeedResults() bool {
@@ -803,6 +809,22 @@ func OutputColumnIsAlwaysNull(e RelExpr, col opt.ColumnID) bool {
 	return false
 }
 
+// IsPartialIndex returns true if the table's index at the given ordinal is
+// a partial index.
+func IsPartialIndex(tabMeta *opt.TableMeta, ord cat.IndexOrdinal) bool {
+	_, isPartial := tabMeta.PartialIndexPredicates[ord]
+	return isPartial
+}
+
+// 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. IsPartialIndex should be used first to
+// determine if the index is a partial index.
+func PartialIndexPredicate(tabMeta *opt.TableMeta, ord cat.IndexOrdinal) FiltersExpr {
+	p := tabMeta.PartialIndexPredicates[ord]
+	return *p.(*FiltersExpr)
+}
+
 // FKCascades stores metadata necessary for building cascading queries.
 type FKCascades []FKCascade
 

pkg/sql/opt/memo/expr_format.go

@@ -1220,6 +1220,7 @@ func FormatPrivate(f *ExprFmtCtx, private interface{}, physProps *physical.Requi
 	case *ScanPrivate:
 		// Don't output name of index if it's the primary index.
 		tab := f.Memo.metadata.Table(t.Table)
+		tabMeta := f.Memo.metadata.TableMeta(t.Table)
 		if t.Index == cat.PrimaryIndex {
 			fmt.Fprintf(f.Buffer, " %s", tableAlias(f, t.Table))
 		} else {
@@ -1228,7 +1229,7 @@ func FormatPrivate(f *ExprFmtCtx, private interface{}, physProps *physical.Requi
 		if ScanIsReverseFn(f.Memo.Metadata(), t, &physProps.Ordering) {
 			f.Buffer.WriteString(",rev")
 		}
-		if _, ok := tab.Index(t.Index).Predicate(); ok {
+		if IsPartialIndex(tabMeta, t.Index) {
 			f.Buffer.WriteString(",partial")
 		}
 

pkg/sql/opt/optbuilder/select.go

@@ -711,9 +711,21 @@ func (b *Builder) addPartialIndexPredicatesForTable(tabMeta *opt.TableMeta) {
 		// Wrap the scalar in a FiltersItem.
 		filter := b.factory.ConstructFiltersItem(scalar)
 
-		// If the expression contains non-immutable operators, do not add it to
-		// the table metadata.
+		// Expressions with non-immutable operators are not supported as partial
+		// index predicates, so add a replacement expression of False. This is
+		// done for two reasons:
+		//
+		//   1. TableMeta.PartialIndexPredicates is a source of truth within the
+		//      optimizer for determining which indexes are partial. It is safer
+		//      to use a False predicate than no predicate so that the optimizer
+		//      won't incorrectly assume that the index is a full index.
+		//   2. A partial index with a False predicate will never be used to
+		//      satisfy a query, effectively making these non-immutable partial
+		//      index predicates not possible to use.
+		//
 		if filter.ScalarProps().VolatilitySet.HasStable() || filter.ScalarProps().VolatilitySet.HasVolatile() {
+			fals := memo.FiltersExpr{b.factory.ConstructFiltersItem(memo.FalseSingleton)}
+			tabMeta.AddPartialIndexPredicate(indexOrd, &fals)
 			return
 		}
 

pkg/sql/opt/xform/custom_funcs.go

@@ -202,8 +202,8 @@ func (c *CustomFuncs) GeneratePartialIndexScans(
 	// Iterate over all partial indexes.
 	iter := makeScanIndexIter(c.e.mem, scanPrivate, rejectNonPartialIndexes)
 	for iter.Next() {
-		pred := tabMeta.PartialIndexPredicates[iter.IndexOrdinal()]
-		remainingFilters, ok := c.im.FiltersImplyPredicate(filters, *pred.(*memo.FiltersExpr))
+		pred := memo.PartialIndexPredicate(tabMeta, iter.IndexOrdinal())
+		remainingFilters, ok := c.im.FiltersImplyPredicate(filters, pred)
 		if !ok {
 			// The filters do not imply the predicate, so the partial index
 			// cannot be used.

pkg/sql/opt/xform/scan_index_iter.go

@@ -144,15 +144,15 @@ func (it *scanIndexIter) Next() bool {
 		}
 
 		if it.hasRejectFlag(rejectPartialIndexes | rejectNonPartialIndexes) {
-			_, ok := it.currIndex.Predicate()
+			isPartialIndex := memo.IsPartialIndex(it.tabMeta, it.indexOrd)
 
 			// Skip over partial indexes if rejectPartialIndexes is set.
-			if it.hasRejectFlag(rejectPartialIndexes) && ok {
+			if it.hasRejectFlag(rejectPartialIndexes) && isPartialIndex {
 				continue
 			}
 
 			// Skip over non-partial indexes if rejectNonPartialIndexes is set.
-			if it.hasRejectFlag(rejectNonPartialIndexes) && !ok {
+			if it.hasRejectFlag(rejectNonPartialIndexes) && !isPartialIndex {
 				continue
 			}
 		}