diff --git a/pkg/sql/opt/memo/expr_format.go b/pkg/sql/opt/memo/expr_format.go index 1a254c0e7b97..b9f781993de7 100644 --- a/pkg/sql/opt/memo/expr_format.go +++ b/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) diff --git a/pkg/sql/opt/optbuilder/select.go b/pkg/sql/opt/optbuilder/select.go index eb1860b1710f..8c33d2785267 100644 --- a/pkg/sql/opt/optbuilder/select.go +++ b/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) { diff --git a/pkg/sql/opt/optbuilder/testdata/select b/pkg/sql/opt/optbuilder/testdata/select index 2ebd5db6b263..9e68d967f9d0 100644 --- a/pkg/sql/opt/optbuilder/testdata/select +++ b/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) diff --git a/pkg/sql/opt/table_meta.go b/pkg/sql/opt/table_meta.go index 545a42c497cb..67b38ea29d99 100644 --- a/pkg/sql/opt/table_meta.go +++ b/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.