Merge #49860

49860: opt: add partial index predicates to TableMeta r=mgartner a=mgartner

With this commit, `optbuilder` now adds partial index predicates of a
table, as a `map[cat.IndexOrdinal]ScalarExpr`, to `TableMeta` when
building SELECT queries. These predicates will be necessary in order to
determine if a partial index can be used to satisfy a query.

Release note: None

Co-authored-by: Marcus Gartner <marcus@cockroachlabs.com>

pkg/sql/opt/memo/expr_format.go

@@ -352,6 +352,19 @@ func (f *ExprFmtCtx) formatRelational(e RelExpr, tp treeprinter.Node) {
  f.formatExpr(tab.ComputedCols[col], c.Child(f.ColumnString(col)))
  }
  }
+ if tab.PartialIndexPredicates != nil {
+ c := tp.Child("partial index predicates")
+ indexOrds := make([]cat.IndexOrdinal, 0, len(tab.PartialIndexPredicates))
+ for ord := range tab.PartialIndexPredicates {
+ indexOrds = append(indexOrds, ord)
+ }
+ sort.Ints(indexOrds)
+ for _, ord := range indexOrds {
+ name := string(tab.Table.Index(ord).Name())
+ f.Buffer.Reset()
+ f.formatScalarWithLabel(name, tab.PartialIndexPredicates[ord], c)
+ }
+ }
  }
  if c := t.Constraint; c != nil {
  if c.IsContradiction() {
@@ -726,6 +739,17 @@ func (f *ExprFmtCtx) formatRelational(e RelExpr, tp treeprinter.Node) {
 }
 
 func (f *ExprFmtCtx) formatScalar(scalar opt.ScalarExpr, tp treeprinter.Node) {
+ f.formatScalarWithLabel("", scalar, tp)
+}
+
+func (f *ExprFmtCtx) formatScalarWithLabel(
+ label string, scalar opt.ScalarExpr, tp treeprinter.Node,
+) {
+ f.Buffer.Reset()
+ if label != "" {
+ f.Buffer.WriteString(label)
+ f.Buffer.WriteString(": ")
+ }
  switch scalar.Op() {
  case opt.ProjectionsOp, opt.AggregationsOp, opt.FKChecksOp, opt.KVOptionsOp:
  // Omit empty lists (except filters).
@@ -741,7 +765,6 @@ func (f *ExprFmtCtx) formatScalar(scalar opt.ScalarExpr, tp treeprinter.Node) {
  }
 
  case opt.IfErrOp:
- f.Buffer.Reset()
  fmt.Fprintf(f.Buffer, "%v", scalar.Op())
  f.FormatScalarProps(scalar)
 
@@ -758,7 +781,6 @@ func (f *ExprFmtCtx) formatScalar(scalar opt.ScalarExpr, tp treeprinter.Node) {
  return
 
  case opt.AggFilterOp:
- f.Buffer.Reset()
  fmt.Fprintf(f.Buffer, "%v", scalar.Op())
  f.FormatScalarProps(scalar)
  tp = tp.Child(f.Buffer.String())
@@ -825,7 +847,6 @@ func (f *ExprFmtCtx) formatScalar(scalar opt.ScalarExpr, tp treeprinter.Node) {
  }
 
  var intercepted bool
- f.Buffer.Reset()
  if f.HasFlags(ExprFmtHideScalars) && ScalarFmtInterceptor != nil {
  if str := ScalarFmtInterceptor(f, scalar); str != "" {
  f.Buffer.WriteString(str)

pkg/sql/opt/optbuilder/select.go

@@ -525,6 +525,7 @@ func (b *Builder) buildScan(
 
  b.addCheckConstraintsForTable(tabMeta)
  b.addComputedColsForTable(tabMeta)
+ b.addPartialIndexPredicatesForTable(tabMeta)
 
  outScope.expr = b.factory.ConstructScan(&private)
 
@@ -615,7 +616,7 @@ func (b *Builder) addComputedColsForTable(tabMeta *opt.TableMeta) {
  }
  expr, err := parser.ParseExpr(tabCol.ComputedExprStr())
  if err != nil {
- continue
+ panic(err)
  }
 
  if tableScope == nil {
@@ -631,6 +632,53 @@ func (b *Builder) addComputedColsForTable(tabMeta *opt.TableMeta) {
  }
 }
 
+// addPartialIndexPredicatesForTable finds all partial indexes in the table and
+// adds their predicates to the table metadata (see
+// TableMeta.PartialIndexPredicates). The predicates are converted from strings
+// to ScalarExprs here.
+func (b *Builder) addPartialIndexPredicatesForTable(tabMeta *opt.TableMeta) {
+ tab := tabMeta.Table
+
+ // Find the first partial index.
+ numIndexes := tab.IndexCount()
+ indexOrd := 0
+ for ; indexOrd < numIndexes; indexOrd++ {
+ if _, ok := tab.Index(indexOrd).Predicate(); ok {
+ break
+ }
+ }
+
+ // Return early if there are no partial indexes. Only partial indexes have
+ // predicates.
+ if indexOrd == numIndexes {
+ return
+ }
+
+ // Create a scope that can be used for building the scalar expressions.
+ tableScope := b.allocScope()
+ tableScope.appendColumnsFromTable(tabMeta, &tabMeta.Alias)
+
+ // Skip to the first partial index we found above.
+ for ; indexOrd < numIndexes; indexOrd++ {
+ index := tab.Index(indexOrd)
+ pred, ok := index.Predicate()
+
+ // If the index is not a partial index, do nothing.
+ if !ok {
+ continue
+ }
+
+ expr, err := parser.ParseExpr(pred)
+ if err != nil {
+ panic(err)
+ }
+
+ texpr := tableScope.resolveAndRequireType(expr, types.Bool)
+ scalar := b.buildScalar(texpr, tableScope, nil, nil, nil)
+ tabMeta.AddPartialIndexPredicate(indexOrd, scalar)
+ }
+}
+
 func (b *Builder) buildSequenceSelect(
  seq cat.Sequence, seqName *tree.TableName, inScope *scope,
 ) (outScope *scope) {

pkg/sql/opt/optbuilder/testdata/select

@@ -1268,3 +1268,26 @@ with &1
  │ └── xyzw.w:7 => w:11
  └── projections
  └── a:1 + x:8 [as="?column?":12]
+
+# Populates table metadata with partial index predicates.
+exec-ddl
+CREATE TABLE partial_index (
+ k INT PRIMARY KEY,
+ u INT,
+ v INT,
+ INDEX u (u) WHERE u = 1,
+ INDEX uv (u, v),
+ INDEX v (v) WHERE v > 100 AND v < 200 AND u > 50
+)
+----
+
+build
+SELECT k FROM partial_index
+----
+project
+ ├── columns: k:1!null
+ └── scan partial_index
+ ├── columns: k:1!null u:2 v:3
+ └── partial index predicates
+ ├── u: u:2 = 1
+ └── v: ((v:3 > 100) AND (v:3 < 200)) AND (u:2 > 50)

pkg/sql/opt/table_meta.go

@@ -149,6 +149,12 @@ type TableMeta struct {
  // more detail.
  ComputedCols map[ColumnID]ScalarExpr
 
+ // PartialIndexPredicates is a map from an index ordinal on the table to
+ // a ScalarExpr representing the predicate on the corresponding partial
+ // index. If an index is not a partial index, it will not have an entry in
+ // the map.
+ PartialIndexPredicates map[cat.IndexOrdinal]ScalarExpr
+
  // anns annotates the table metadata with arbitrary data.
  anns [maxTableAnnIDCount]interface{}
 }
@@ -202,6 +208,15 @@ func (tm *TableMeta) AddComputedCol(colID ColumnID, computedCol ScalarExpr) {
  tm.ComputedCols[colID] = computedCol
 }
 
+// AddPartialIndexPredicate adds a partial index predicate to the table's
+// metadata.
+func (tm *TableMeta) AddPartialIndexPredicate(ord cat.IndexOrdinal, pred ScalarExpr) {
+ if tm.PartialIndexPredicates == nil {
+ tm.PartialIndexPredicates = make(map[cat.IndexOrdinal]ScalarExpr)
+ }
+ tm.PartialIndexPredicates[ord] = pred
+}
+
 // TableAnnotation returns the given annotation that is associated with the
 // given table. If the table has no such annotation, TableAnnotation returns
 // nil.