opt: propagate set operation output types to input columns

This commit updates the optbuilder logic for set operations in which
the types of the input columns do not match the types of the output
columns. This can happen if a column on one side has type Unknown,
but the corresponding column on the other side has a known type such
as Int. The known type must be propagated to the side with the unknown
type to prevent errors in the execution engine related to decoding
types.

If there are any column types on either side that don't match the output,
then the optbuilder propagates the output types of the set operation down
to the input columns by wrapping the side with mismatched types in a
Project operation. The Project operation passes through columns that
already have the correct type, and creates cast expressions for those
that don't.

Fixes #34524

Release note (bug fix): Fixed an error that happened when executing
some set operations containing only nulls in one of the input columns.

pkg/sql/opt/exec/execbuilder/testdata/union

@@ -1,4 +1,4 @@
-# LogicTest: local-opt
+# LogicTest: local-opt fakedist-opt
 
 statement ok
 CREATE TABLE uniontest (
@@ -101,3 +101,12 @@ render                 ·              ·                 ("?column?", "?column?
 ·                      row 0, expr 0  ''                ·                                     ·
 ·                      row 0, expr 1  ''                ·                                     ·
 ·                      row 0, expr 2  'x'               ·                                     ·
+
+statement ok
+CREATE TABLE a (a INT PRIMARY KEY)
+
+# Regression test for #34524. This test is here because the issue still exists
+# in the heuristic planner.
+query I
+(SELECT NULL FROM a) EXCEPT (VALUES((SELECT 1 FROM a LIMIT 1)), (1))
+----

pkg/sql/opt/norm/custom_funcs.go

@@ -658,32 +658,6 @@ func (c *CustomFuncs) AreProjectionsCorrelated(
 	return false
 }
 
-// ProjectColMapLeft returns a Projections operator that maps the left side
-// columns in a SetPrivate to the output columns in it. Useful for replacing set
-// operations with simpler constructs.
-func (c *CustomFuncs) ProjectColMapLeft(set *memo.SetPrivate) memo.ProjectionsExpr {
-	return c.projectColMapSide(set.OutCols, set.LeftCols)
-}
-
-// ProjectColMapRight returns a Project operator that maps the right side
-// columns in a SetPrivate to the output columns in it. Useful for replacing set
-// operations with simpler constructs.
-func (c *CustomFuncs) ProjectColMapRight(set *memo.SetPrivate) memo.ProjectionsExpr {
-	return c.projectColMapSide(set.OutCols, set.RightCols)
-}
-
-// projectColMapSide implements the side-agnostic logic from ProjectColMapLeft
-// and ProjectColMapRight.
-func (c *CustomFuncs) projectColMapSide(toList, fromList opt.ColList) memo.ProjectionsExpr {
-	items := make(memo.ProjectionsExpr, len(toList))
-	for idx, fromCol := range fromList {
-		toCol := toList[idx]
-		items[idx].Element = c.f.ConstructVariable(fromCol)
-		items[idx].Col = toCol
-	}
-	return items
-}
-
 // MakeEmptyColSet returns a column set with no columns in it.
 func (c *CustomFuncs) MakeEmptyColSet() opt.ColSet {
 	return opt.ColSet{}
@@ -885,6 +859,39 @@ func (c *CustomFuncs) ZipOuterCols(zip memo.ZipExpr) opt.ColSet {
 	return colSet
 }
 
+// ----------------------------------------------------------------------
+//
+// Set Rules
+//   Custom match and replace functions used with set.opt rules.
+//
+// ----------------------------------------------------------------------
+
+// ProjectColMapLeft returns a Projections operator that maps the left side
+// columns in a SetPrivate to the output columns in it. Useful for replacing set
+// operations with simpler constructs.
+func (c *CustomFuncs) ProjectColMapLeft(set *memo.SetPrivate) memo.ProjectionsExpr {
+	return c.projectColMapSide(set.OutCols, set.LeftCols)
+}
+
+// ProjectColMapRight returns a Project operator that maps the right side
+// columns in a SetPrivate to the output columns in it. Useful for replacing set
+// operations with simpler constructs.
+func (c *CustomFuncs) ProjectColMapRight(set *memo.SetPrivate) memo.ProjectionsExpr {
+	return c.projectColMapSide(set.OutCols, set.RightCols)
+}
+
+// projectColMapSide implements the side-agnostic logic from ProjectColMapLeft
+// and ProjectColMapRight.
+func (c *CustomFuncs) projectColMapSide(toList, fromList opt.ColList) memo.ProjectionsExpr {
+	items := make(memo.ProjectionsExpr, len(toList))
+	for idx, fromCol := range fromList {
+		toCol := toList[idx]
+		items[idx].Element = c.f.ConstructVariable(fromCol)
+		items[idx].Col = toCol
+	}
+	return items
+}
+
 // ----------------------------------------------------------------------
 //
 // Boolean Rules

pkg/sql/opt/norm/rules/select.opt

@@ -344,33 +344,3 @@
     $input
     (RemoveFiltersItem $filters $item)
 )
-
-# EliminateUnionAllLeft replaces a union all with a right side having a
-# cardinality of zero, with just the left side operand.
-[EliminateUnionAllLeft, Normalize]
-(UnionAll
-    $left:*
-    $right:* & (HasZeroRows $right)
-    $colmap:*
-)
-=>
-(Project
-    $left
-    (ProjectColMapLeft $colmap)
-    (MakeEmptyColSet)
-)
-
-# EliminateUnionAllRight replaces a union all with a left side having a
-# cardinality of zero, with just the right side operand.
-[EliminateUnionAllRight, Normalize]
-(UnionAll
-    $left:* & (HasZeroRows $left)
-    $right:*
-    $colmap:*
-)
-=>
-(Project
-    $right
-    (ProjectColMapRight $colmap)
-    (MakeEmptyColSet)
-)

pkg/sql/opt/norm/rules/set.opt

@@ -0,0 +1,33 @@
+# =============================================================================
+# set.opt contains normalization rules for set operators.
+# =============================================================================
+
+# EliminateUnionAllLeft replaces a union all with a right side having a
+# cardinality of zero, with just the left side operand.
+[EliminateUnionAllLeft, Normalize]
+(UnionAll
+    $left:*
+    $right:* & (HasZeroRows $right)
+    $colmap:*
+)
+=>
+(Project
+    $left
+    (ProjectColMapLeft $colmap)
+    (MakeEmptyColSet)
+)
+
+# EliminateUnionAllRight replaces a union all with a left side having a
+# cardinality of zero, with just the right side operand.
+[EliminateUnionAllRight, Normalize]
+(UnionAll
+    $left:* & (HasZeroRows $left)
+    $right:*
+    $colmap:*
+)
+=>
+(Project
+    $right
+    (ProjectColMapRight $colmap)
+    (MakeEmptyColSet)
+)

pkg/sql/opt/norm/testdata/rules/select

@@ -1133,51 +1133,3 @@ values
  ├── cardinality: [0 - 0]
  ├── key: ()
  └── fd: ()-->(1)
-
-# --------------------------------------------------
-# EliminateUnionAllLeft
-# --------------------------------------------------
-
-opt expect=EliminateUnionAllLeft
-SELECT k FROM
-  (SELECT k FROM b)
-  UNION ALL
-  (SELECT k FROM b WHERE k IN ())
-----
-project
- ├── columns: k:11(int)
- ├── scan b
- │    ├── columns: b.k:1(int!null)
- │    └── key: (1)
- └── projections
-      └── variable: b.k [type=int, outer=(1)]
-
-# --------------------------------------------------
-# EliminateUnionAllRight
-# --------------------------------------------------
-
-opt expect=EliminateUnionAllRight
-SELECT k FROM
-  (SELECT k FROM b WHERE Null)
-  UNION ALL
-  (SELECT k FROM b)
-----
-project
- ├── columns: k:11(int)
- ├── scan b
- │    ├── columns: b.k:6(int!null)
- │    └── key: (6)
- └── projections
-      └── variable: b.k [type=int, outer=(6)]
-
-opt
-SELECT k FROM
-  (SELECT k FROM b WHERE False)
-  UNION ALL
-  (SELECT k FROM b WHERE i IN ())
-----
-values
- ├── columns: k:11(int)
- ├── cardinality: [0 - 0]
- ├── key: ()
- └── fd: ()-->(11)

pkg/sql/opt/norm/testdata/rules/set

@@ -0,0 +1,59 @@
+exec-ddl
+CREATE TABLE b (k INT PRIMARY KEY, i INT, f FLOAT, s STRING NOT NULL, j JSON)
+----
+TABLE b
+ ├── k int not null
+ ├── i int
+ ├── f float
+ ├── s string not null
+ ├── j jsonb
+ └── INDEX primary
+      └── k int not null
+
+# --------------------------------------------------
+# EliminateUnionAllLeft
+# --------------------------------------------------
+
+opt expect=EliminateUnionAllLeft
+SELECT k FROM
+  (SELECT k FROM b)
+  UNION ALL
+  (SELECT k FROM b WHERE k IN ())
+----
+project
+ ├── columns: k:11(int)
+ ├── scan b
+ │    ├── columns: b.k:1(int!null)
+ │    └── key: (1)
+ └── projections
+      └── variable: b.k [type=int, outer=(1)]
+
+# --------------------------------------------------
+# EliminateUnionAllRight
+# --------------------------------------------------
+
+opt expect=EliminateUnionAllRight
+SELECT k FROM
+  (SELECT k FROM b WHERE Null)
+  UNION ALL
+  (SELECT k FROM b)
+----
+project
+ ├── columns: k:11(int)
+ ├── scan b
+ │    ├── columns: b.k:6(int!null)
+ │    └── key: (6)
+ └── projections
+      └── variable: b.k [type=int, outer=(6)]
+
+opt
+SELECT k FROM
+  (SELECT k FROM b WHERE False)
+  UNION ALL
+  (SELECT k FROM b WHERE i IN ())
+----
+values
+ ├── columns: k:11(int)
+ ├── cardinality: [0 - 0]
+ ├── key: ()
+ └── fd: ()-->(11)

pkg/sql/opt/ops/relational.opt

@@ -605,6 +605,9 @@ define Union {
 #   Left:  [2, 1]
 #   Right: [4, 3]
 #   Out:   [5, 6]  <-- synthesized output columns
+#
+# To make normalization rules and execution simpler, both inputs to the set op
+# must have matching types.
 [Private]
 define SetPrivate {
     LeftCols  ColList

pkg/sql/opt/optbuilder/scope.go

@@ -135,29 +135,38 @@ func (s *scope) replace() *scope {
 }
 
 // appendColumnsFromScope adds newly bound variables to this scope.
-// The groups in the new columns are reset to 0.
+// The expressions in the new columns are reset to nil.
 func (s *scope) appendColumnsFromScope(src *scope) {
 	l := len(s.cols)
 	s.cols = append(s.cols, src.cols...)
-	// We want to reset the groups, as these become pass-through columns in the
-	// new scope.
+	// We want to reset the expressions, as these become pass-through columns in
+	// the new scope.
 	for i := l; i < len(s.cols); i++ {
 		s.cols[i].scalar = nil
 	}
 }
 
 // appendColumns adds newly bound variables to this scope.
-// The groups in the new columns are reset to 0.
+// The expressions in the new columns are reset to nil.
 func (s *scope) appendColumns(cols []scopeColumn) {
 	l := len(s.cols)
 	s.cols = append(s.cols, cols...)
-	// We want to reset the groups, as these become pass-through columns in the
-	// new scope.
+	// We want to reset the expressions, as these become pass-through columns in
+	// the new scope.
 	for i := l; i < len(s.cols); i++ {
 		s.cols[i].scalar = nil
 	}
 }
 
+// appendColumn adds a newly bound variable to this scope.
+// The expression in the new column is reset to nil.
+func (s *scope) appendColumn(col *scopeColumn) {
+	s.cols = append(s.cols, *col)
+	// We want to reset the expression, as this becomes a pass-through column in
+	// the new scope.
+	s.cols[len(s.cols)-1].scalar = nil
+}
+
 // addExtraColumns adds the given columns as extra columns, ignoring any
 // duplicate columns that are already in the scope.
 func (s *scope) addExtraColumns(cols []scopeColumn) {