opt: normalize partial index PUT/DEL projections to false

The `SimplifyPartialIndexProjections` normalization rule has been added
that normalizes synthesized partial index PUT and DEL columns to False
when it is guaranteed that a mutation will not require changes to the
associated partial index. This normalization can lead to further
normalizations, such as pruning columns that the synthesized projections
relied on.

The motivation for this change is to allow fully disjoint updates to
different columns in the same row, when the columns are split across
different families. By pruning columns not needed to maintain a partial
index, we're not forced to scan all column families. This can ultimately
reduce contention during updates.

Release note (performance improvement): UPDATE operations on tables with
partial indexes no longer evaluate partial index predicate expressions
when it is guaranteed that the operation will not alter the state of the
partial index. In some cases, this can eliminate fetching the existing
value of columns that are referenced in partial index predicates.

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

pkg/sql/opt/norm/prune_cols_funcs.go

@@ -44,6 +44,15 @@ func (c *CustomFuncs) NeededExplainCols(private *memo.ExplainPrivate) opt.ColSet
 // in turn trigger the PruneMutationInputCols rule.
 func (c *CustomFuncs) NeededMutationCols(
 	private *memo.MutationPrivate, uniqueChecks memo.UniqueChecksExpr, fkChecks memo.FKChecksExpr,
+) opt.ColSet {
+	return c.neededMutationCols(private, uniqueChecks, fkChecks, true /* includePartialIndexCols */)
+}
+
+func (c *CustomFuncs) neededMutationCols(
+	private *memo.MutationPrivate,
+	uniqueChecks memo.UniqueChecksExpr,
+	fkChecks memo.FKChecksExpr,
+	includePartialIndexCols bool,
 ) opt.ColSet {
 	var cols opt.ColSet
 
@@ -60,8 +69,10 @@ func (c *CustomFuncs) NeededMutationCols(
 	addCols(private.FetchCols)
 	addCols(private.UpdateCols)
 	addCols(private.CheckCols)
-	addCols(private.PartialIndexPutCols)
-	addCols(private.PartialIndexDelCols)
+	if includePartialIndexCols {
+		addCols(private.PartialIndexPutCols)
+		addCols(private.PartialIndexDelCols)
+	}
 	addCols(private.ReturnCols)
 	addCols(opt.OptionalColList(private.PassthroughCols))
 	if private.CanaryCol != 0 {

pkg/sql/opt/norm/rules/mutation.opt

@@ -0,0 +1,38 @@
+# =============================================================================
+# mutation.opt contains normalization rules for the mutation operators.
+# =============================================================================
+
+# SimplifyPartialIndexProjections converts partial index PUT and DEL projected
+# expressions to false when it is guaranteed that the mutation will not require
+# changes to the associated partial index. 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.
+[SimplifyPartialIndexProjections, Normalize]
+(Update
+    $project:(Project $input:* $projections:* $passthrough:*)
+    $uniqueChecks:*
+    $fkChecks:*
+    $mutationPrivate:* &
+        ^(ColsAreEmpty
+            $simplifiableCols:(SimplifiablePartialIndexProjectCols
+                $mutationPrivate
+                $uniqueChecks
+                $fkChecks
+                $projections
+            )
+        )
+)
+=>
+(Update
+    (Project
+        $input
+        (SimplifyPartialIndexProjections
+            $projections
+            $simplifiableCols
+        )
+        $passthrough
+    )
+    $uniqueChecks
+    $fkChecks
+    $mutationPrivate
+)