opt: add helper functions for partial indexes and scans

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
 			}
 		}