opt: prune partial index columns and simplify partial index projections

pkg/sql/backfill/backfill.go

@@ -253,9 +253,10 @@ func (cb *ColumnBackfiller) RunColumnBackfillChunk(
 ) (roachpb.Key, error) {
  // TODO(dan): Tighten up the bound on the requestedCols parameter to
  // makeRowUpdater.
- requestedCols := make([]descpb.ColumnDescriptor, 0, len(tableDesc.Columns)+len(cb.added))
+ requestedCols := make([]descpb.ColumnDescriptor, 0, len(tableDesc.Columns)+len(cb.added)+len(cb.dropped))
  requestedCols = append(requestedCols, tableDesc.Columns...)
  requestedCols = append(requestedCols, cb.added...)
+ requestedCols = append(requestedCols, cb.dropped...)
  ru, err := row.MakeUpdater(
  ctx,
  txn,

pkg/sql/logictest/testdata/logic_test/partial_index

@@ -1281,6 +1281,54 @@ SELECT * FROM inv_c@i WHERE j @> '{"x": "y"}' AND s IN ('foo', 'bar') ORDER BY k
 1 {"num": 1, "x": "y"} foo
 3 {"num": 3, "x": "y"} bar
 
+# Updates and Upserts with fetch columns pruned.
+
+statement ok
+CREATE TABLE prune (
+ a INT PRIMARY KEY,
+ b INT,
+ c INT,
+ d INT,
+ INDEX idx (b) WHERE c > 0,
+ FAMILY (a),
+ FAMILY (b),
+ FAMILY (c),
+ FAMILY (d)
+)
+
+statement ok
+INSERT INTO prune (a, b, c, d) VALUES (1, 2, 3, 4)
+
+# Test that an update is successful when fetch columns b and c are pruned
+# because an update to idx is not required.
+statement ok
+UPDATE prune SET d = d + 1 WHERE a = 1
+
+query IIII rowsort
+SELECT * FROM prune@idx WHERE c > 0
+----
+1 2 3 5
+
+# Test that an upsert is successful when fetch columns b and c are pruned
+# because an update to idx is not required.
+statement ok
+UPSERT INTO prune (a, d) VALUES (1, 6)
+
+query IIII rowsort
+SELECT * FROM prune@idx WHERE c > 0
+----
+1 2 3 6
+
+# Test that an upsert is successful when fetch columns b and c are pruned
+# because an update to idx is not required.
+statement ok
+INSERT INTO prune (a, d) VALUES (1, 6) ON CONFLICT (a) DO UPDATE SET d = 7
+
+query IIII rowsort
+SELECT * FROM prune@idx WHERE c > 0
+----
+1 2 3 7
+
 # Regression tests for #52318. Mutations on partial indexes in the
 # DELETE_AND_WRITE_ONLY state should update the indexes correctly.
 subtest regression_52318

pkg/sql/opt/memo/expr.go

@@ -658,7 +658,7 @@ func (s *ScanPrivate) IsLocking() bool {
 // index, nil is returned.
 func (s *ScanPrivate) PartialIndexPredicate(md *opt.Metadata) FiltersExpr {
  tabMeta := md.TableMeta(s.Table)
- p, ok := tabMeta.PartialIndexPredicates[s.Index]
+ p, ok := tabMeta.PartialIndexPredicate(s.Index)
  if !ok {
  // The index is not a partial index.
  return nil

pkg/sql/opt/memo/expr_format.go

@@ -352,17 +352,18 @@ func (f *ExprFmtCtx) formatRelational(e RelExpr, tp treeprinter.Node) {
  f.formatExpr(tab.ComputedCols[col], c.Child(f.ColumnString(col)))
  }
  }
- if tab.PartialIndexPredicates != nil {
+ partialIndexPredicates := tab.PartialIndexPredicatesForFormattingOnly()
+ if partialIndexPredicates != nil {
  c := tp.Child("partial index predicates")
- indexOrds := make([]cat.IndexOrdinal, 0, len(tab.PartialIndexPredicates))
- for ord := range tab.PartialIndexPredicates {
+ indexOrds := make([]cat.IndexOrdinal, 0, len(partialIndexPredicates))
+ for ord := range 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)
+ f.formatScalarWithLabel(name, partialIndexPredicates[ord], c)
  }
  }
  }

pkg/sql/opt/memo/testdata/logprops/delete

@@ -64,7 +64,6 @@ project
  ├── volatile, mutations
  ├── key: (5)
  ├── fd: ()-->(1), (5)-->(2-4)
- ├── prune: (1-4)
  └── select
  ├── columns: a:8(int!null) b:9(int) c:10(int!null) d:11(int) rowid:12(int!null) e:13(int) crdb_internal_mvcc_timestamp:14(decimal)
  ├── key: (12)
@@ -107,7 +106,6 @@ project
  ├── volatile, mutations
  ├── key: ()
  ├── fd: ()-->(1-5)
- ├── prune: (1-4)
  └── select
  ├── columns: a:8(int!null) b:9(int) c:10(int!null) d:11(int) rowid:12(int!null) e:13(int) crdb_internal_mvcc_timestamp:14(decimal)
  ├── cardinality: [0 - 1]
@@ -148,7 +146,6 @@ project
  ├── volatile, mutations
  ├── key: (5)
  ├── fd: (2)==(3), (3)==(2), (5)-->(1-4)
- ├── prune: (1-4)
  └── select
  ├── columns: a:8(int!null) b:9(int!null) c:10(int!null) d:11(int) rowid:12(int!null) e:13(int) crdb_internal_mvcc_timestamp:14(decimal)
  ├── key: (12)

pkg/sql/opt/metadata.go

@@ -396,7 +396,7 @@ func (md *Metadata) AddTable(tab cat.Table, alias *tree.TableName) TableID {
 // ScalarExpr to new column IDs. It takes as arguments a ScalarExpr and a
 // mapping of old column IDs to new column IDs, and returns a new ScalarExpr.
 // This function is used when duplicating Constraints, ComputedCols, and
-// PartialIndexPredicates. DuplicateTable requires this callback function,
+// partialIndexPredicates. DuplicateTable requires this callback function,
 // rather than performing the remapping itself, because remapping column IDs
 // requires constructing new expressions with norm.Factory. The norm package
 // depends on opt, and cannot be imported here.
@@ -450,9 +450,9 @@ func (md *Metadata) DuplicateTable(
  // Create new partial index predicate expressions by remapping the column
  // IDs in each ScalarExpr.
  var partialIndexPredicates map[cat.IndexOrdinal]ScalarExpr
- if len(tabMeta.PartialIndexPredicates) > 0 {
- partialIndexPredicates = make(map[cat.IndexOrdinal]ScalarExpr, len(tabMeta.PartialIndexPredicates))
- for idxOrd, e := range tabMeta.PartialIndexPredicates {
+ if len(tabMeta.partialIndexPredicates) > 0 {
+ partialIndexPredicates = make(map[cat.IndexOrdinal]ScalarExpr, len(tabMeta.partialIndexPredicates))
+ for idxOrd, e := range tabMeta.partialIndexPredicates {
  partialIndexPredicates[idxOrd] = remapColumnIDs(e, colMap)
  }
  }
@@ -464,7 +464,7 @@ func (md *Metadata) DuplicateTable(
  IgnoreForeignKeys: tabMeta.IgnoreForeignKeys,
  Constraints: constraints,
  ComputedCols: computedCols,
- PartialIndexPredicates: partialIndexPredicates,
+ partialIndexPredicates: partialIndexPredicates,
  })
 
  return newTabID

pkg/sql/opt/metadata_test.go

@@ -281,7 +281,7 @@ func TestIndexColumns(t *testing.T) {
 // TestDuplicateTable tests that we can extract a set of columns from an index ordinal.
 func TestDuplicateTable(t *testing.T) {
  cat := testcat.New()
- _, err := cat.ExecuteDDL("CREATE TABLE a (b BOOL, b2 BOOL)")
+ _, err := cat.ExecuteDDL("CREATE TABLE a (b BOOL, b2 BOOL, INDEX (b2) WHERE b)")
  if err != nil {
  t.Fatal(err)
  }
@@ -332,11 +332,12 @@ func TestDuplicateTable(t *testing.T) {
  t.Errorf("expected computed column to reference new column ID %d, got %d", dupB, col)
  }
 
- if dupTabMeta.PartialIndexPredicates == nil || dupTabMeta.PartialIndexPredicates[1] == nil {
+ pred, isPartialIndex := dupTabMeta.PartialIndexPredicate(1)
+ if !isPartialIndex {
  t.Fatalf("expected partial index predicates to be duplicated")
  }
 
- col = dupTabMeta.PartialIndexPredicates[1].(*memo.VariableExpr).Col
+ col = pred.(*memo.VariableExpr).Col
  if col == b {
  t.Errorf("expected partial index predicate to reference new column ID %d, got %d", dupB, col)
  }

pkg/sql/opt/norm/BUILD.bazel

@@ -14,6 +14,7 @@ go_library(
  "join_funcs.go",
  "limit_funcs.go",
  "list_sorter.go",
+ "mutation_funcs.go",
  "ordering_funcs.go",
  "project_builder.go",
  "project_funcs.go",

pkg/sql/opt/norm/mutation_funcs.go

@@ -0,0 +1,119 @@
+// Copyright 2020 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package norm
+
+import (
+ "github.com/cockroachdb/cockroach/pkg/sql/opt"
+ "github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
+)
+
+// SimplifiablePartialIndexProjectCols returns the set of projected partial
+// index PUT and DEL columns with expressions that can be simplified to false.
+// These projected expressions can only be simplified to false when an UPDATE
+// mutates neither the associated index's columns nor the columns referenced in
+// the partial index predicate.
+func (c *CustomFuncs) SimplifiablePartialIndexProjectCols(
+ private *memo.MutationPrivate,
+ uniqueChecks memo.UniqueChecksExpr,
+ fkChecks memo.FKChecksExpr,
+ projections memo.ProjectionsExpr,
+) opt.ColSet {
+ tabMeta := c.mem.Metadata().TableMeta(private.Table)
+
+ // Determine the set of target table columns that need to be updated. Notice
+ // that we collect the target table column IDs, not the update column IDs.
+ var updateCols opt.ColSet
+ for ord, col := range private.UpdateCols {
+ if col != 0 {
+ updateCols.Add(tabMeta.MetaID.ColumnID(ord))
+ }
+ }
+
+ // Determine the set of columns needed for the mutation operator, excluding
+ // the partial index PUT and DEL columns.
+ neededMutationCols := c.neededMutationCols(private, uniqueChecks, fkChecks, false /* includePartialIndexCols */)
+
+ // Determine the set of project columns that are already simplified to
+ // false.
+ var simplifiedProjectCols opt.ColSet
+ for i := range projections {
+ project := &projections[i]
+ if project.Element == memo.FalseSingleton {
+ simplifiedProjectCols.Add(project.Col)
+ }
+ }
+
+ // Columns that are required by the mutation operator and columns that
+ // have already been simplified to false are ineligible to be simplified.
+ ineligibleCols := neededMutationCols.Union(simplifiedProjectCols)
+
+ // ord is an ordinal into the mutation's PartialIndexPutCols and
+ // PartialIndexDelCols, which both have entries for each partial index
+ // defined on the table.
+ ord := -1
+ var cols opt.ColSet
+ for i, n := 0, tabMeta.Table.DeletableIndexCount(); i < n; i++ {
+ pred, isPartialIndex := tabMeta.PartialIndexPredicate(i)
+
+ // Skip non-partial indexes.
+ if !isPartialIndex {
+ continue
+ }
+ ord++
+
+ // If the columns being updated are part of the index or referenced in
+ // the partial index predicate, then updates to the index may be
+ // required. Therefore, the partial index PUT and DEL columns cannot be
+ // simplified.
+ //
+ // Note that we use the set of index columns where the virtual
+ // columns have been mapped to their source columns. Virtual columns
+ // are never part of the updated columns. Updates to source columns
+ // trigger index changes.
+ predFilters := *pred.(*memo.FiltersExpr)
+ indexAndPredCols := tabMeta.IndexColumnsMapVirtual(i)
+ indexAndPredCols.UnionWith(predFilters.OuterCols())
+ if indexAndPredCols.Intersects(updateCols) {
+ continue
+ }
+
+ // Add the projected PUT column if it is eligible to be simplified.
+ putCol := private.PartialIndexPutCols[ord]
+ if !ineligibleCols.Contains(putCol) {
+ cols.Add(putCol)
+ }
+
+ // Add the projected DEL column if it is eligible to be simplified.
+ delCol := private.PartialIndexDelCols[ord]
+ if !ineligibleCols.Contains(delCol) {
+ cols.Add(delCol)
+ }
+ }
+
+ return cols
+}
+
+// SimplifyPartialIndexProjections returns a new projection expression with any
+// projected column's expression simplified to false if the column exists in
+// simplifiableCols.
+func (c *CustomFuncs) SimplifyPartialIndexProjections(
+ projections memo.ProjectionsExpr, simplifiableCols opt.ColSet,
+) memo.ProjectionsExpr {
+ simplified := make(memo.ProjectionsExpr, len(projections))
+ for i := range projections {
+ if col := projections[i].Col; simplifiableCols.Contains(col) {
+ simplified[i] = c.f.ConstructProjectionsItem(memo.FalseSingleton, col)
+ } else {
+ simplified[i] = projections[i]
+ }
+ }
+ return simplified
+}