Merge pull request #98115 from rytaft/backport22.2-98042

release-22.2: opt: fix incorrect results returned by non-recursive WITH RECURSIVE CTE (#98115)

pkg/sql/logictest/testdata/logic_test/with

@@ -1141,3 +1141,36 @@ SELECT i, j FROM u@u_j_idx
 query TTTTTTTT
 EXPERIMENTAL SCRUB TABLE u WITH OPTIONS INDEX ALL
 ----
+
+# Regression tests for #93370. Do not convert a non-recursive CTE
+# that uses UNION ALL and WITH RECURSIVE to UNION.
+query I rowsort
+WITH RECURSIVE
+   x(id) AS
+     (SELECT 1 UNION ALL SELECT id+1 FROM x WHERE id < 3 ),
+   y(id) AS
+     (SELECT * FROM x UNION ALL SELECT * FROM x)
+ SELECT * FROM y
+----
+1
+2
+3
+1
+2
+3
+
+statement ok
+CREATE TABLE t93370 (i INT);
+INSERT INTO t93370 VALUES (1), (2), (3)
+
+query I rowsort
+WITH RECURSIVE
+   y(id) AS (SELECT * FROM t93370 UNION ALL SELECT * FROM t93370)
+ SELECT * FROM y
+----
+1
+2
+3
+1
+2
+3

pkg/sql/opt/exec/execbuilder/testdata/with

@@ -231,3 +231,66 @@ vectorized: true
 │
 └── • values
       size: 1 column, 3 rows
+
+# Regression tests for #93370. Do not convert a non-recursive CTE
+# that uses UNION ALL and WITH RECURSIVE to UNION.
+query T
+EXPLAIN WITH RECURSIVE
+   x(id) AS
+     (SELECT 1 UNION ALL SELECT id+1 FROM x WHERE id < 3 ),
+   y(id) AS
+     (SELECT * FROM x UNION ALL SELECT * FROM x)
+ SELECT * FROM y
+----
+distribution: local
+vectorized: true
+·
+• root
+│
+├── • render
+│   │
+│   └── • union all
+│       │
+│       ├── • scan buffer
+│       │     label: buffer 2 (x)
+│       │
+│       └── • scan buffer
+│             label: buffer 2 (x)
+│
+└── • subquery
+    │ id: @S1
+    │ original sql: SELECT 1 UNION ALL SELECT id + 1 FROM x WHERE id < 3
+    │ exec mode: all rows
+    │
+    └── • buffer
+        │ label: buffer 2 (x)
+        │
+        └── • recursive cte
+            │
+            └── • values
+                  size: 1 column, 1 row
+
+statement ok
+CREATE TABLE t93370 (i INT)
+
+query T
+EXPLAIN WITH RECURSIVE
+   y(id) AS (SELECT * FROM t93370 UNION ALL SELECT * FROM t93370)
+ SELECT * FROM y
+----
+distribution: local
+vectorized: true
+·
+• render
+│
+└── • union all
+    │
+    ├── • scan
+    │     missing stats
+    │     table: t93370@t93370_pkey
+    │     spans: FULL SCAN
+    │
+    └── • scan
+          missing stats
+          table: t93370@t93370_pkey
+          spans: FULL SCAN

pkg/sql/opt/optbuilder/testdata/with

@@ -1149,7 +1149,7 @@ WITH RECURSIVE cte(a, b) AS (
 ----
 with &2 (cte)
  ├── columns: a:9!null b:10!null
- ├── union
+ ├── union-all
  │    ├── columns: "?column?":7!null "?column?":8!null
  │    ├── left columns: "?column?":1 "?column?":2
  │    ├── right columns: "?column?":5 "?column?":6
@@ -2035,3 +2035,87 @@ with &4 (qux)
            │                   │    └── filters (true)
            │                   └── filters (true)
            └── filters (true)
+
+# Regression tests for #93370. Do not convert a non-recursive CTE
+# that uses UNION ALL and WITH RECURSIVE to UNION.
+build
+WITH RECURSIVE
+   x(id) AS
+     (SELECT 1 UNION ALL SELECT id+1 FROM x WHERE id < 3 ),
+   y(id) AS
+     (SELECT * FROM x UNION ALL SELECT * FROM x)
+ SELECT * FROM y
+----
+with &2 (x)
+ ├── columns: id:9
+ ├── recursive-c-t-e
+ │    ├── columns: id:2
+ │    ├── working table binding: &1
+ │    ├── initial columns: "?column?":1
+ │    ├── recursive columns: "?column?":4
+ │    ├── fake-rel
+ │    │    └── columns: id:2
+ │    ├── project
+ │    │    ├── columns: "?column?":1!null
+ │    │    ├── values
+ │    │    │    └── ()
+ │    │    └── projections
+ │    │         └── 1 [as="?column?":1]
+ │    └── project
+ │         ├── columns: "?column?":4!null
+ │         ├── select
+ │         │    ├── columns: id:3!null
+ │         │    ├── with-scan &1 (x)
+ │         │    │    ├── columns: id:3
+ │         │    │    └── mapping:
+ │         │    │         └──  id:2 => id:3
+ │         │    └── filters
+ │         │         └── id:3 < 3
+ │         └── projections
+ │              └── id:3 + 1 [as="?column?":4]
+ └── with &4 (y)
+      ├── columns: id:9
+      ├── union-all
+      │    ├── columns: id:8
+      │    ├── left columns: id:5
+      │    ├── right columns: id:7
+      │    ├── with-scan &2 (x)
+      │    │    ├── columns: id:5
+      │    │    └── mapping:
+      │    │         └──  id:2 => id:5
+      │    └── with-scan &2 (x)
+      │         ├── columns: id:7
+      │         └── mapping:
+      │              └──  id:2 => id:7
+      └── with-scan &4 (y)
+           ├── columns: id:9
+           └── mapping:
+                └──  id:8 => id:9
+
+exec-ddl
+CREATE TABLE t93370 (i INT)
+----
+
+build
+WITH RECURSIVE
+   y(id) AS (SELECT * FROM t93370 UNION ALL SELECT * FROM t93370)
+ SELECT * FROM y
+----
+with &2 (y)
+ ├── columns: id:11
+ ├── union-all
+ │    ├── columns: i:10
+ │    ├── left columns: t93370.i:1
+ │    ├── right columns: t93370.i:6
+ │    ├── project
+ │    │    ├── columns: t93370.i:1
+ │    │    └── scan t93370
+ │    │         └── columns: t93370.i:1 rowid:2!null crdb_internal_mvcc_timestamp:3 tableoid:4
+ │    └── project
+ │         ├── columns: t93370.i:6
+ │         └── scan t93370
+ │              └── columns: t93370.i:6 rowid:7!null crdb_internal_mvcc_timestamp:8 tableoid:9
+ └── with-scan &2 (y)
+      ├── columns: id:11
+      └── mapping:
+           └──  i:10 => id:11

pkg/sql/opt/optbuilder/with.go

@@ -315,7 +315,7 @@ func (b *Builder) buildCTE(
 	recursiveScope := b.buildStmt(recursive, initialTypes /* desiredTypes */, cteScope)
 	if numRefs == 0 {
 		// Build this as a non-recursive CTE.
-		cteScope := b.buildSetOp(tree.UnionOp, false /* all */, inScope, initialScope, recursiveScope)
+		cteScope := b.buildSetOp(tree.UnionOp, isUnionAll, inScope, initialScope, recursiveScope)
 		return cteScope.expr, b.getCTECols(cteScope, cte.Name), nil
 	}