Merge cockroachdb#45597

45597: opt: extend WithUses and improve NeededMutationCols r=RaduBerinde a=RaduBerinde

This change extends the WithUses rule prop to keep track of the columns actually
used by `WithScan`s.

This set is then used by NeededMutationCols to make sure we never prune mutation
input columns that are used by FK checks. Currently this never happens, but
because of fairly subtle reasons: FKs happen to require indexes on both sides,
and those indexes force the mutation operator to require values for those
columns. This won't be the case anymore with upsert FK checks, for which this
fix is required.

The new information will also be used (in a future change) to prune unused
columns of `With` itself.

Release note: None

Co-authored-by: Radu Berinde <radu@cockroachlabs.com>

pkg/sql/opt/exec/execbuilder/mutation.go

@@ -40,6 +40,18 @@ func (b *Builder) buildMutationInput(
 		return execPlan{}, err
 	}
 
+	if p.WithID != 0 {
+		// The input might have extra columns that are used only by FK checks; make
+		// sure we don't project them away.
+		cols := inputExpr.Relational().OutputCols.Copy()
+		for _, c := range colList {
+			cols.Remove(c)
+		}
+		for c, ok := cols.Next(0); ok; c, ok = cols.Next(c + 1) {
+			colList = append(colList, c)
+		}
+	}
+
 	input, err = b.ensureColumns(input, colList, nil, inputExpr.ProvidedPhysical().Ordering)
 	if err != nil {
 		return execPlan{}, err

pkg/sql/opt/memo/expr_format.go

@@ -662,9 +662,19 @@ func (f *ExprFmtCtx) formatRelational(e RelExpr, tp treeprinter.Node) {
 		if r.JoinSize > 1 {
 			tp.Childf("join-size: %d", r.JoinSize)
 		}
-		if len(relational.Shared.Rule.WithUses) > 0 {
-			// Go map
-			tp.Childf("cte-uses: %v", relational.Shared.Rule.WithUses)
+		if withUses := relational.Shared.Rule.WithUses; len(withUses) > 0 {
+			n := tp.Childf("cte-uses")
+			ids := make([]opt.WithID, 0, len(withUses))
+			for id := range withUses {
+				ids = append(ids, id)
+			}
+			sort.Slice(ids, func(i, j int) bool {
+				return ids[i] < ids[j]
+			})
+			for _, id := range ids {
+				info := withUses[id]
+				n.Childf("&%d: count=%d used-columns=%s", id, info.Count, info.UsedCols)
+			}
 		}
 	}
 

pkg/sql/opt/memo/testdata/logprops/with

@@ -24,7 +24,8 @@ with &1 (foo)
       ├── key: (3)
       ├── fd: (3)-->(4)
       ├── prune: (3,4)
-      └── cte-uses: map[1:1]
+      └── cte-uses
+           └── &1: count=1 used-columns=(1,2)
 
 # Side effects should be propagated up to the top-level from the Binding side
 # of a WITH.
@@ -93,7 +94,8 @@ with &1 (foo)
       ├── key: ()
       ├── fd: ()-->(3)
       ├── prune: (3)
-      ├── cte-uses: map[1:1]
+      ├── cte-uses
+      │    └── &1: count=1 used-columns=(1)
       ├── with-scan &1 (foo)
       │    ├── columns: "?column?":2(int!null)
       │    ├── mapping:
@@ -102,7 +104,8 @@ with &1 (foo)
       │    ├── key: ()
       │    ├── fd: ()-->(2)
       │    ├── prune: (2)
-      │    └── cte-uses: map[1:1]
+      │    └── cte-uses
+      │         └── &1: count=1 used-columns=(1)
       └── projections
            └── div [as="?column?":3, type=decimal, side-effects]
                 ├── const: 1 [type=int]
@@ -138,7 +141,8 @@ with &1 (foo)
       ├── key: ()
       ├── fd: ()-->(3)
       ├── prune: (3)
-      ├── cte-uses: map[1:1]
+      ├── cte-uses
+      │    └── &1: count=1 used-columns=(1)
       ├── with-scan &1 (foo)
       │    ├── columns: int8:2(int)
       │    ├── mapping:
@@ -147,7 +151,8 @@ with &1 (foo)
       │    ├── key: ()
       │    ├── fd: ()-->(2)
       │    ├── prune: (2)
-      │    └── cte-uses: map[1:1]
+      │    └── cte-uses
+      │         └── &1: count=1 used-columns=(1)
       └── projections
            └── const: 1 [as="?column?":3, type=int]
 

pkg/sql/opt/norm/custom_funcs.go

@@ -2045,7 +2045,7 @@ func (c *CustomFuncs) CanInlineWith(binding, expr memo.RelExpr, private *memo.Wi
 	if binding.Relational().CanHaveSideEffects {
 		return false
 	}
-	return c.WithUses(expr)[private.ID] <= 1
+	return c.WithUses(expr)[private.ID].Count <= 1
 }
 
 // InlineWith replaces all references to the With expression in input (via
@@ -2079,9 +2079,8 @@ func (c *CustomFuncs) InlineWith(binding, input memo.RelExpr, priv *memo.WithPri
 	return replace(input).(memo.RelExpr)
 }
 
-// WithUses returns a map mapping WithIDs to the number of times a given With
-// expression is referenced in the given expression.
-func (c *CustomFuncs) WithUses(r opt.Expr) map[opt.WithID]int {
+// WithUses returns the WithUsesMap for the given expression.
+func (c *CustomFuncs) WithUses(r opt.Expr) props.WithUsesMap {
 	switch e := r.(type) {
 	case memo.RelExpr:
 		relProps := e.Relational()
@@ -2106,29 +2105,46 @@ func (c *CustomFuncs) WithUses(r opt.Expr) map[opt.WithID]int {
 	}
 }
 
-// deriveWithUses computes the number of times each WithScan is referenced. It's
-// used to decide if we can inline a With or not.
-func (c *CustomFuncs) deriveWithUses(r opt.Expr) map[opt.WithID]int {
-	var result map[opt.WithID]int
+// deriveWithUses collects information about WithScans in the expression.
+func (c *CustomFuncs) deriveWithUses(r opt.Expr) props.WithUsesMap {
+	// We don't allow the information to escape the scope of the WITH itself, so
+	// we exclude that ID from the results.
+	var excludedID opt.WithID
+
 	switch e := r.(type) {
 	case *memo.WithScanExpr:
-		result = map[opt.WithID]int{e.With: 1}
+		info := props.WithUseInfo{
+			Count:    1,
+			UsedCols: e.InCols.ToSet(),
+		}
+		return props.WithUsesMap{e.With: info}
+
+	case *memo.WithExpr:
+		excludedID = e.ID
+
 	default:
-		for i, n := 0, r.ChildCount(); i < n; i++ {
-			for id, useCount := range c.WithUses(r.Child(i)) {
-				if result == nil {
-					result = make(map[opt.WithID]int)
-				}
-				result[id] = result[id] + useCount
-			}
+		if opt.IsMutationOp(e) {
+			// Note: this can still be 0.
+			excludedID = e.Private().(*memo.MutationPrivate).WithID
 		}
 	}
 
-	// Don't allow the use count to escape the scope of the WITH itself.
-	if w, ok := r.(*memo.WithExpr); ok {
-		delete(result, w.ID)
+	var result props.WithUsesMap
+	for i, n := 0, r.ChildCount(); i < n; i++ {
+		childUses := c.WithUses(r.Child(i))
+		for id, info := range childUses {
+			if id == excludedID {
+				continue
+			}
+			if result == nil {
+				result = make(props.WithUsesMap, len(childUses))
+			}
+			existing := result[id]
+			existing.Count += info.Count
+			existing.UsedCols.UnionWith(info.UsedCols)
+			result[id] = existing
+		}
 	}
-
 	return result
 }
 

pkg/sql/opt/norm/prune_cols.go

@@ -42,7 +42,9 @@ func (c *CustomFuncs) NeededExplainCols(private *memo.ExplainPrivate) opt.ColSet
 // needed by the mutation private; it simply returns the input columns that are
 // referenced by it. Other rules filter the FetchCols, CheckCols, etc. and can
 // in turn trigger the PruneMutationInputCols rule.
-func (c *CustomFuncs) NeededMutationCols(private *memo.MutationPrivate) opt.ColSet {
+func (c *CustomFuncs) NeededMutationCols(
+	private *memo.MutationPrivate, checks memo.FKChecksExpr,
+) opt.ColSet {
 	var cols opt.ColSet
 
 	// Add all input columns referenced by the mutation private.
@@ -64,6 +66,13 @@ func (c *CustomFuncs) NeededMutationCols(private *memo.MutationPrivate) opt.ColS
 		cols.Add(private.CanaryCol)
 	}
 
+	if private.WithID != 0 {
+		for i := range checks {
+			withUses := c.WithUses(checks[i].Check)
+			cols.UnionWith(withUses[private.WithID].UsedCols)
+		}
+	}
+
 	return cols
 }
 

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

@@ -481,7 +481,7 @@
     $input:*
     $checks:*
     $mutationPrivate:* &
-        (CanPruneCols $input $needed:(NeededMutationCols $mutationPrivate))
+        (CanPruneCols $input $needed:(NeededMutationCols $mutationPrivate $checks))
 )
 =>
 ((OpName)

pkg/sql/opt/norm/testdata/rules/prune_cols

@@ -2870,7 +2870,8 @@ with &1 (foo)
       ├── stats: [rows=1000]
       ├── cost: 0.01
       ├── prune: (6)
-      └── cte-uses: map[1:1]
+      └── cte-uses
+           └── &1: count=1 used-columns=(2)
 
 # --------------------------------------------------
 # PruneUnionAllCols