opt: fetch minimal set of columns on returning mutations

Previously, we used to fetch all columns when a mutation contained
a `RETURNING` clause. This is an issue because it forces us to
retrieve unnecessary data and creates extra contention.
This change adds logic to compute the minimal set of required columns
and fetches only those.

Fixes cockroachdb#30618.
Unblocks cockroachdb#30624.

Release note: None

pkg/sql/opt/exec/execbuilder/testdata/delete

@@ -142,18 +142,15 @@ count           ·         ·
 ·               spans     ALL
 ·               limit     10
 
-# TODO(andyk): Prune columns so that index-join is not necessary.
 query TTT
 EXPLAIN DELETE FROM indexed WHERE value = 5 LIMIT 10 RETURNING id
 ----
-render                     ·         ·
- └── run                   ·         ·
-      └── delete           ·         ·
-           │               from      indexed
-           │               strategy  deleter
-           └── index-join  ·         ·
-                │          table     indexed@primary
-                └── scan   ·         ·
-·                          table     indexed@indexed_value_idx
-·                          spans     /5-/6
-·                          limit     10
+render               ·         ·
+ └── run             ·         ·
+      └── delete     ·         ·
+           │         from      indexed
+           │         strategy  deleter
+           └── scan  ·         ·
+·                    table     indexed@indexed_value_idx
+·                    spans     /5-/6
+·                    limit     10

pkg/sql/opt/exec/execbuilder/testdata/orderby

@@ -502,7 +502,7 @@ render               ·         ·          (b)        ·
       └── delete     ·         ·          (a, b, c)  ·
            │         from      t          ·          ·
            │         strategy  deleter    ·          ·
-           └── scan  ·         ·          (a, b, c)  ·
+           └── scan  ·         ·          (a, b)     ·
 ·                    table     t@primary  ·          ·
 ·                    spans     /3-/3/#    ·          ·
 

pkg/sql/opt/norm/prune_cols.go

@@ -13,6 +13,7 @@ package norm
 import (
 	"github.com/cockroachdb/cockroach/pkg/sql/opt"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/exec"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/props"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/props/physical"
@@ -509,3 +510,56 @@ func DerivePruneCols(e memo.RelExpr) opt.ColSet {
 
 	return relProps.Rule.PruneCols
 }
+
+// CanPruneMutationReturnCols checks whether the mutations return columns can
+// be pruned. This is the base case for the PruneMutationReturnCols rule.
+func (c *CustomFuncs) CanPruneMutationReturnCols(private *memo.MutationPrivate) bool {
+	if private.ReturnCols == nil || private.ReturnPruned {
+		return false
+	}
+
+	return true
+}
+
+// PruneMutationReturnCols rewrites the given mutation private to no longer
+// keep ReturnCols that are not referenced by the RETURNING clause or are not
+// part of the primary key. The caller must have already done the analysis to
+// prove that these columns are not needed, by calling CanPruneMutationReturnCols.
+func (c *CustomFuncs) PruneMutationReturnCols(
+	private *memo.MutationPrivate, projections memo.ProjectionsExpr, passthrough opt.ColSet,
+) *memo.MutationPrivate {
+	newPrivate := *private
+	newReturnCols := make(opt.ColList, len(private.ReturnCols))
+
+	// Find all the columns referenced by the projections.
+	returningOrdSet := exec.ColumnOrdinalSet{}
+	for _, projection := range projections {
+		outCols := projection.ScalarProps(c.mem).OuterCols
+		outCols.ForEach(func(i opt.ColumnID) {
+			returningOrdSet.Add(private.Table.ColumnOrdinal(i))
+		})
+	}
+	passthrough.ForEach(func(i opt.ColumnID) {
+		if i > 0 {
+			returningOrdSet.Add(private.Table.ColumnOrdinal(i))
+		}
+	})
+
+	// The columns of the primary index are always returned regardless of
+	// whether they are referenced.
+	tab := c.mem.Metadata().Table(private.Table)
+	primaryIndex := tab.Index(0)
+	for i, n := 0, primaryIndex.KeyColumnCount(); i < n; i++ {
+		returningOrdSet.Add(primaryIndex.Column(i).Ordinal)
+	}
+
+	for i := 0; i < len(private.ReturnCols); i++ {
+		if returningOrdSet.Contains(i) {
+			newReturnCols[i] = private.ReturnCols[i]
+		}
+	}
+
+	newPrivate.ReturnCols = newReturnCols
+	newPrivate.ReturnPruned = true
+	return &newPrivate
+}

pkg/sql/opt/norm/rules/prune_cols.opt

@@ -462,3 +462,27 @@
     $checks
     $mutationPrivate
 )
+
+# PruneReturningCols removes columns from the mutation operator's ReturnCols
+# set if they are not used in the RETURNING clause of the mutation.
+# Removing ReturnCols will then allow the PruneMutationFetchCols to be more
+# conservative with the fetch columns.
+[PruneMutationReturnCols, Normalize]
+(Project
+    $input: (Insert | Update | Upsert | Delete
+        $innerInput:*
+        $checks:*
+        $mutationPrivate:* & (CanPruneMutationReturnCols $mutationPrivate))
+    $projections:*
+    $passthrough:*
+)
+=>
+(Project
+    ((OpName $input)
+        $innerInput
+        $checks
+        (PruneMutationReturnCols $mutationPrivate $projections $passthrough))
+    $projections
+    $passthrough
+)
+