opt: prune update/upsert fetch columns not needed for partial indexes

Indexed columns of partial indexes are now only fetched for UPDATE and
UPSERT operations when needed. They are pruned in cases where it is
guaranteed that they are not needed to build old or new index entries.
For example, consider the table and UPDATE:

    CREATE TABLE t (
      a INT PRIMARY KEY,
      b INT,
      c INT,
      d INT,
      INDEX (b) WHERE c > 0,
      FAMILY (a), FAMILY (b), FAMILY (c), FAMILY (d)
    )

    UPDATE t SET d = d + 1 WHERE a = 1

The partial index is guaranteed not to change with this UPDATE because
neither its indexed columns nor the columns referenced in its predicate
are mutating. Therefore, the existing values of b do not need to be
fetched to maintain the state of the partial index. Furthermore, the
primary index does require the existing values of b because no columns
in b's family are mutating. So, b can be pruned from the UPDATE's fetch
columns.

Release note (performance improvement): Previously, indexed columns of
partial indexes were always fetched for UPDATEs and UPSERTs. Now they
are only fetched if they are required for maintaining the state of the
index. If an UPDATE or UPSERT mutates columns that are neither indexed by a
partial index nor referenced in a partial index predicate, they will no
longer be fetched (assuming that they are not needed to maintain the
state of other indexes, including the primary index).

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/norm/prune_cols_funcs.go

@@ -144,26 +144,26 @@ func (c *CustomFuncs) NeededMutationFetchCols(
 
 		// Make sure to consider indexes that are being added or dropped.
 		for i, n := 0, tabMeta.Table.DeletableIndexCount(); i < n; i++ {
-			// If the columns being updated are not part of the index and the
-			// index is not a partial index, then the update does not require
-			// changes to the index. Partial indexes may be updated (even when a
-			// column in the index is not changing) when rows that were not
-			// previously in the index must be added to the index because they
-			// now satisfy the partial index predicate.
+			// If the columns being updated are not part of the index, then the
+			// update does not require changes to the index. Partial indexes may
+			// be updated (even when a column in the index is not changing) when
+			// the predicate references columns that are being updated. For
+			// example, rows that were not previously in the index must be added
+			// to the index because they now satisfy the partial index
+			// predicate, requiring the index columns to be fetched.
 			//
 			// 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.
-			//
-			// TODO(mgartner): Index columns are not necessary when neither the
-			// index columns nor the columns referenced in the partial index
-			// predicate are being updated. We should prune mutation fetch
-			// columns when this is the case, rather than always marking index
-			// columns of partial indexes as "needed".
 			indexCols := tabMeta.IndexColumnsMapVirtual(i)
-			_, isPartialIndex := tabMeta.Table.Index(i).Predicate()
-			if !indexCols.Intersects(updateCols) && !isPartialIndex {
+			pred, isPartialIndex := tabMeta.PartialIndexPredicate(i)
+			indexAndPredCols := indexCols.Copy()
+			if isPartialIndex {
+				predFilters := *pred.(*memo.FiltersExpr)
+				indexAndPredCols.UnionWith(predFilters.OuterCols())
+			}
+			if !indexAndPredCols.Intersects(updateCols) {
 				continue
 			}