Merge #43161

43161: opt: disambiguate WithScan expressions using a unique ID. r=savoie a=savoie

Fixes #43148.

As was previously the case with Values expressions (see #35004), it is
possible for multiple identical WithScan expressions to appear within
the same tree (this occurs if they are scanning the same CTE and have no
columns after normalization). This might cause a number of subtle bugs
with anything depending on expressions having unique hashes, such as
defining different required physical properties for each identical
`WithScan`. This patch adds a memo-unique identifier to each
`WithScanExpr` to force distinguishability.

Release note (bug fix): Previously, the optimizer could panic in a
specific situation where it would prune all the columns of multiple
scans of the same CTE and then try to define different required physical
properties for each scan. This seems to have been a possible bug since
the addition of multi-use CTEs in 19.2, but is hard to trigger without
the not-yet-released LimitHint physical property. This patch makes all
CTE scans uniquely identifiable even after column pruning.

Co-authored-by: Céline O'Neil <[email protected]>

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

@@ -1505,20 +1505,20 @@ func (b *Builder) buildRecursiveCTE(rec *memo.RecursiveCTEExpr) (execPlan, error
 }
 
 func (b *Builder) buildWithScan(withScan *memo.WithScanExpr) (execPlan, error) {
-	id := withScan.ID
+	withID := withScan.With
 	var e *builtWithExpr
 	for i := range b.withExprs {
-		if b.withExprs[i].id == id {
+		if b.withExprs[i].id == withID {
 			e = &b.withExprs[i]
 			break
 		}
 	}
 	if e == nil {
-		panic(errors.AssertionFailedf("couldn't find With expression with ID %d", id))
+		panic(errors.AssertionFailedf("couldn't find With expression with ID %d", withID))
 	}
 
 	var label bytes.Buffer
-	fmt.Fprintf(&label, "buffer %d", withScan.ID)
+	fmt.Fprintf(&label, "buffer %d", withScan.With)
 	if withScan.Name != "" {
 		fmt.Fprintf(&label, " (%s)", withScan.Name)
 	}

pkg/sql/opt/memo/expr_format.go

@@ -215,7 +215,7 @@ func (f *ExprFmtCtx) formatRelational(e RelExpr, tp treeprinter.Node) {
 		}
 
 	case *WithScanExpr:
-		fmt.Fprintf(f.Buffer, "%v &%d", e.Op(), t.ID)
+		fmt.Fprintf(f.Buffer, "%v &%d", e.Op(), t.With)
 		if t.Name != "" {
 			fmt.Fprintf(f.Buffer, " (%s)", t.Name)
 		}

pkg/sql/opt/memo/interner.go

@@ -463,7 +463,7 @@ func (h *hasher) HashSequenceID(val opt.SequenceID) {
 	h.HashUint64(uint64(val))
 }
 
-func (h *hasher) HashValuesID(val opt.ValuesID) {
+func (h *hasher) HashUniqueID(val opt.UniqueID) {
 	h.HashUint64(uint64(val))
 }
 
@@ -790,7 +790,7 @@ func (h *hasher) IsSequenceIDEqual(l, r opt.SequenceID) bool {
 	return l == r
 }
 
-func (h *hasher) IsValuesIDEqual(l, r opt.ValuesID) bool {
+func (h *hasher) IsUniqueIDEqual(l, r opt.UniqueID) bool {
 	return l == r
 }
 

pkg/sql/opt/metadata.go

@@ -77,8 +77,8 @@ type Metadata struct {
 	// sequences stores information about each metadata sequence, indexed by SequenceID.
 	sequences []cat.Sequence
 
-	// values is the highest id for a Values clause that has been assigned.
-	values ValuesID
+	// currUniqueID is the highest UniqueID that has been assigned.
+	currUniqueID UniqueID
 
 	// deps stores information about all data source objects depended on by the
 	// query, as well as the privileges required to access them. The objects are
@@ -446,17 +446,19 @@ func (md *Metadata) AllSequences() []cat.Sequence {
 	return md.sequences
 }
 
-// ValuesID uniquely identifies the usage of a values clause within the scope of a
-// query.
+// UniqueID should be used to disambiguate multiple uses of an expression
+// within the scope of a query. For example, a UniqueID field should be
+// added to an expression type if two instances of that type might otherwise
+// be indistinguishable based on the values of their other fields.
 //
 // See the comment for Metadata for more details on identifiers.
-type ValuesID uint64
+type UniqueID uint64
 
-// NextValuesID returns a fresh ValuesID which is guaranteed to never have been
-// allocated prior in this memo.
-func (md *Metadata) NextValuesID() ValuesID {
-	md.values++
-	return md.values
+// NextUniqueID returns a fresh UniqueID which is guaranteed to never have been
+// previously allocated in this memo.
+func (md *Metadata) NextUniqueID() UniqueID {
+	md.currUniqueID++
+	return md.currUniqueID
 }
 
 // AddView adds a new reference to a view used by the query.

pkg/sql/opt/norm/custom_funcs.go

@@ -1166,7 +1166,7 @@ func (c *CustomFuncs) ConstructEmptyValues(cols opt.ColSet) memo.RelExpr {
 	}
 	return c.f.ConstructValues(memo.EmptyScalarListExpr, &memo.ValuesPrivate{
 		Cols: colList,
-		ID:   c.mem.Metadata().NextValuesID(),
+		ID:   c.mem.Metadata().NextUniqueID(),
 	})
 }
 
@@ -1951,7 +1951,7 @@ func (c *CustomFuncs) InlineWith(binding, input memo.RelExpr, priv *memo.WithPri
 	replace = func(nd opt.Expr) opt.Expr {
 		switch t := nd.(type) {
 		case *memo.WithScanExpr:
-			if t.ID == priv.ID {
+			if t.With == priv.ID {
 				// TODO(justin): it might be worth carefully walking the tree and
 				// renaming variables as we do this replacement so that this projection
 				// is unnecessary (assuming there's at most one reference to the
@@ -2008,7 +2008,7 @@ func (c *CustomFuncs) deriveWithUses(r opt.Expr) map[opt.WithID]int {
 	var result map[opt.WithID]int
 	switch e := r.(type) {
 	case *memo.WithScanExpr:
-		result = map[opt.WithID]int{e.ID: 1}
+		result = map[opt.WithID]int{e.With: 1}
 	default:
 		for i, n := 0, r.ChildCount(); i < n; i++ {
 			for id, useCount := range c.WithUses(r.Child(i)) {

pkg/sql/opt/norm/decorrelate.go

@@ -262,7 +262,7 @@ func (c *CustomFuncs) HoistValuesSubquery(
 
 	values := c.f.ConstructValues(newRows, &memo.ValuesPrivate{
 		Cols: private.Cols,
-		ID:   c.f.Metadata().NextValuesID(),
+		ID:   c.f.Metadata().NextUniqueID(),
 	})
 	join := c.f.ConstructInnerJoinApply(hoister.input(), values, memo.TrueFilter, memo.EmptyJoinPrivate)
 	outCols := values.Relational().OutputCols
@@ -690,7 +690,7 @@ func (c *CustomFuncs) ConstructBinary(op opt.Operator, left, right opt.ScalarExp
 func (c *CustomFuncs) ConstructNoColsRow() memo.RelExpr {
 	return c.f.ConstructValues(memo.ScalarListWithEmptyTuple, &memo.ValuesPrivate{
 		Cols: opt.ColList{},
-		ID:   c.f.Metadata().NextValuesID(),
+		ID:   c.f.Metadata().NextUniqueID(),
 	})
 }
 

pkg/sql/opt/norm/factory.go

@@ -284,7 +284,7 @@ func (f *Factory) onConstructScalar(scalar opt.ScalarExpr) opt.ScalarExpr {
 func (f *Factory) ConstructZeroValues() memo.RelExpr {
 	return f.ConstructValues(memo.EmptyScalarListExpr, &memo.ValuesPrivate{
 		Cols: opt.ColList{},
-		ID:   f.Metadata().NextValuesID(),
+		ID:   f.Metadata().NextUniqueID(),
 	})
 }
 

pkg/sql/opt/norm/factory_test.go

@@ -47,7 +47,7 @@ func TestSimplifyFilters(t *testing.T) {
 	// Filters expression evaluates to False if any operand is False.
 	vals := f.ConstructValues(memo.ScalarListWithEmptyTuple, &memo.ValuesPrivate{
 		Cols: opt.ColList{},
-		ID:   f.Metadata().NextValuesID(),
+		ID:   f.Metadata().NextUniqueID(),
 	})
 	filters := memo.FiltersExpr{{Condition: eq}, {Condition: memo.FalseSingleton}, {Condition: eq}}
 	sel := f.ConstructSelect(vals, filters)

pkg/sql/opt/norm/testdata/rules/with

@@ -498,3 +498,48 @@ values
  ├── key: ()
  ├── fd: ()-->(5)
  └── (12,) [type=tuple{int}]
+
+# Regression test for #43148: WithScans with no columns should still be
+# uniquely identifiable. Without this uniqueness, they can't be assigned
+# different required physical properties.
+norm
+WITH cte AS (SELECT * FROM a) (SELECT 1 FROM cte LIMIT 9) UNION (SELECT 1 FROM cte LIMIT 10)
+----
+with &1 (cte)
+ ├── columns: "?column?":18(int!null)
+ ├── cardinality: [0 - 19]
+ ├── key: (18)
+ ├── scan a
+ │    ├── columns: a.k:1(int!null) a.i:2(int) a.f:3(float) a.s:4(string) a.j:5(jsonb)
+ │    ├── key: (1)
+ │    └── fd: (1)-->(2-5)
+ └── union
+      ├── columns: "?column?":18(int!null)
+      ├── left columns: "?column?":11(int)
+      ├── right columns: "?column?":17(int)
+      ├── cardinality: [0 - 19]
+      ├── key: (18)
+      ├── project
+      │    ├── columns: "?column?":11(int!null)
+      │    ├── cardinality: [0 - 9]
+      │    ├── fd: ()-->(11)
+      │    ├── limit
+      │    │    ├── cardinality: [0 - 9]
+      │    │    ├── with-scan &1 (cte)
+      │    │    │    ├── mapping:
+      │    │    │    └── limit hint: 9.00
+      │    │    └── const: 9 [type=int]
+      │    └── projections
+      │         └── const: 1 [type=int]
+      └── project
+           ├── columns: "?column?":17(int!null)
+           ├── cardinality: [0 - 10]
+           ├── fd: ()-->(17)
+           ├── limit
+           │    ├── cardinality: [0 - 10]
+           │    ├── with-scan &1 (cte)
+           │    │    ├── mapping:
+           │    │    └── limit hint: 10.00
+           │    └── const: 10 [type=int]
+           └── projections
+                └── const: 1 [type=int]

pkg/sql/opt/ops/relational.opt

@@ -142,7 +142,7 @@ define ValuesPrivate {
     # Values expressions which appear in different places in the query. In most
     # cases the column set is sufficient to do this, but various rules make it
     # possible to construct Values expressions with no columns.
-    ID ValuesID
+    ID UniqueID
 }
 
 # Select filters rows from its input result set, based on the boolean filter
@@ -845,7 +845,8 @@ define WithScan {
 
 [Private]
 define WithScanPrivate {
-    ID WithID
+    # With identifies the CTE to scan.
+    With WithID
 
     # BindingProps stores the relational properties of the referenced expression.
     BindingProps RelPropsPtr
@@ -864,6 +865,12 @@ define WithScanPrivate {
     # the same tree, so we maintain a mapping from the columns returned from
     # the referenced expression to the referencing expression.
     OutCols ColList
+
+    # ID is a memo-unique identifier which distinguishes between identical
+    # WithScan expressions which appear in different places in the query. In
+    # most cases the column set is sufficient to do this, but various rules make
+    # it possible to construct WithScan expressions with no columns.
+    ID UniqueID
 }
 
 # RecursiveCTE implements the logic of a recursive CTE:

pkg/sql/opt/optbuilder/insert.go

@@ -518,7 +518,7 @@ func (mb *mutationBuilder) buildInputForInsert(inScope *scope, inputRows *tree.S
 		mb.outScope = inScope.push()
 		mb.outScope.expr = mb.b.factory.ConstructValues(memo.ScalarListWithEmptyTuple, &memo.ValuesPrivate{
 			Cols: opt.ColList{},
-			ID:   mb.md.NextValuesID(),
+			ID:   mb.md.NextUniqueID(),
 		})
 		return
 	}

pkg/sql/opt/optbuilder/mutation_builder.go

@@ -1351,10 +1351,11 @@ func (mb *mutationBuilder) projectOrdinals(
 		inCols[i] = mb.outScope.cols[ords[i]].id
 	}
 	out := mb.b.factory.ConstructWithScan(&memo.WithScanPrivate{
-		ID:           mb.withID,
+		With:         mb.withID,
 		InCols:       inCols,
 		OutCols:      outCols,
 		BindingProps: mb.outScope.expr.Relational(),
+		ID:           mb.b.factory.Metadata().NextUniqueID(),
 	})
 	return out, outCols
 }
@@ -1372,10 +1373,11 @@ func (mb *mutationBuilder) makeFKInputScan(
 		outCols[i] = mb.md.AddColumn(c.Alias, c.Type)
 	}
 	scan = mb.b.factory.ConstructWithScan(&memo.WithScanPrivate{
-		ID:           mb.withID,
+		With:         mb.withID,
 		InCols:       inputCols,
 		OutCols:      outCols,
 		BindingProps: mb.outScope.expr.Relational(),
+		ID:           mb.b.factory.Metadata().NextUniqueID(),
 	})
 	return scan, outCols
 }

pkg/sql/opt/optbuilder/select.go

@@ -88,11 +88,12 @@ func (b *Builder) buildDataSource(
 			}
 
 			outScope.expr = b.factory.ConstructWithScan(&memo.WithScanPrivate{
-				ID:           cte.id,
+				With:         cte.id,
 				Name:         string(cte.name.Alias),
 				InCols:       inCols,
 				OutCols:      outCols,
 				BindingProps: cte.bindingProps,
+				ID:           b.factory.Metadata().NextUniqueID(),
 			})
 
 			return outScope
@@ -170,11 +171,12 @@ func (b *Builder) buildDataSource(
 		}
 
 		outScope.expr = b.factory.ConstructWithScan(&memo.WithScanPrivate{
-			ID:           cte.id,
+			With:         cte.id,
 			Name:         string(cte.name.Alias),
 			InCols:       inCols,
 			OutCols:      outCols,
 			BindingProps: cte.bindingProps,
+			ID:           b.factory.Metadata().NextUniqueID(),
 		})
 
 		return outScope
@@ -903,7 +905,7 @@ func (b *Builder) buildFrom(from tree.From, inScope *scope) (outScope *scope) {
 		outScope = inScope.push()
 		outScope.expr = b.factory.ConstructValues(memo.ScalarListWithEmptyTuple, &memo.ValuesPrivate{
 			Cols: opt.ColList{},
-			ID:   b.factory.Metadata().NextValuesID(),
+			ID:   b.factory.Metadata().NextUniqueID(),
 		})
 	}
 

pkg/sql/opt/optbuilder/srfs.go

@@ -121,7 +121,7 @@ func (b *Builder) buildZip(exprs tree.Exprs, inScope *scope) (outScope *scope) {
 	// Construct the zip as a ProjectSet with empty input.
 	input := b.factory.ConstructValues(memo.ScalarListWithEmptyTuple, &memo.ValuesPrivate{
 		Cols: opt.ColList{},
-		ID:   b.factory.Metadata().NextValuesID(),
+		ID:   b.factory.Metadata().NextUniqueID(),
 	})
 	outScope.expr = b.factory.ConstructProjectSet(input, zip)
 	if len(outScope.cols) == 1 {

pkg/sql/opt/optbuilder/values.go

@@ -121,7 +121,7 @@ func (b *Builder) buildValuesClause(
 	colList := colsToColList(outScope.cols)
 	outScope.expr = b.factory.ConstructValues(tuples, &memo.ValuesPrivate{
 		Cols: colList,
-		ID:   b.factory.Metadata().NextValuesID(),
+		ID:   b.factory.Metadata().NextUniqueID(),
 	})
 	return outScope
 }

pkg/sql/opt/optgen/cmd/optgen/metadata.go

@@ -149,7 +149,7 @@ func newMetadata(compiled *lang.CompiledExpr, pkg string) *metadata {
 		"TableID":        {fullName: "opt.TableID", passByVal: true},
 		"SchemaID":       {fullName: "opt.SchemaID", passByVal: true},
 		"SequenceID":     {fullName: "opt.SequenceID", passByVal: true},
-		"ValuesID":       {fullName: "opt.ValuesID", passByVal: true},
+		"UniqueID":       {fullName: "opt.UniqueID", passByVal: true},
 		"WithID":         {fullName: "opt.WithID", passByVal: true},
 		"Ordering":       {fullName: "opt.Ordering", passByVal: true},
 		"OrderingChoice": {fullName: "physical.OrderingChoice", passByVal: true},