fix: the output columns of cte related columns lack the columns required by cte itself (#17576)

* fix: the output columns of cte related columns lack the columns required by cte itself

Signed-off-by: Kould <[email protected]>

* fix: the output columns of cte related columns lack the columns required by cte itself

Signed-off-by: Kould <[email protected]>

* chore: codefmt

Signed-off-by: Kould <[email protected]>

---------

Signed-off-by: Kould <[email protected]>

Author: KKould

src/query/sql/src/executor/physical_plans/physical_union_all.rs

@@ -62,29 +62,45 @@ impl PhysicalPlanBuilder {
         // 1. Prune unused Columns.
         let metadata = self.metadata.read().clone();
         let lazy_columns = metadata.lazy_columns();
-        required.extend(lazy_columns.clone());
-
-        let indices: Vec<usize> = (0..union_all.left_outputs.len())
-            .filter(|index| required.contains(&union_all.output_indexes[*index]))
-            .collect();
-
-        let (left_required, right_required) = if indices.is_empty() {
-            (
-                HashSet::from([union_all.left_outputs[0].0]),
-                HashSet::from([union_all.right_outputs[0].0]),
-            )
+        required.extend(lazy_columns);
+
+        // if the union has a CTE, the output columns are not filtered
+        // otherwise, if the output columns of the union do not contain the columns used by the plan in the union, the expression will fail to obtain data.
+        let (left_required, right_required) = if !union_all.cte_scan_names.is_empty() {
+            let left: ColumnSet = union_all
+                .left_outputs
+                .iter()
+                .map(|(index, _)| *index)
+                .collect();
+            let right: ColumnSet = union_all
+                .right_outputs
+                .iter()
+                .map(|(index, _)| *index)
+                .collect();
+
+            (left, right)
         } else {
-            indices.iter().fold(
+            let indices: Vec<usize> = (0..union_all.left_outputs.len())
+                .filter(|index| required.contains(&union_all.output_indexes[*index]))
+                .collect();
+            if indices.is_empty() {
                 (
-                    HashSet::with_capacity(indices.len()),
-                    HashSet::with_capacity(indices.len()),
-                ),
-                |(mut left, mut right), &index| {
-                    left.insert(union_all.left_outputs[index].0);
-                    right.insert(union_all.right_outputs[index].0);
-                    (left, right)
-                },
-            )
+                    HashSet::from([union_all.left_outputs[0].0]),
+                    HashSet::from([union_all.right_outputs[0].0]),
+                )
+            } else {
+                indices.iter().fold(
+                    (
+                        HashSet::with_capacity(indices.len()),
+                        HashSet::with_capacity(indices.len()),
+                    ),
+                    |(mut left, mut right), &index| {
+                        left.insert(union_all.left_outputs[index].0);
+                        right.insert(union_all.right_outputs[index].0);
+                        (left, right)
+                    },
+                )
+            }
         };
 
         // 2. Build physical plan.

src/query/sql/src/planner/binder/util.rs

@@ -63,12 +63,7 @@ impl Binder {
             RelOperator::RecursiveCteScan(plan) => {
                 cte_scan_names.push(plan.table_name.clone());
                 if cte_types.is_empty() {
-                    cte_types.extend(
-                        plan.fields
-                            .iter()
-                            .map(|f| f.data_type().clone())
-                            .collect::<Vec<DataType>>(),
-                    );
+                    cte_types.extend(plan.fields.iter().map(|f| f.data_type().clone()));
                 }
             }
 

tests/sqllogictests/suites/query/cte/cte.test

@@ -578,10 +578,10 @@ with t(tt) as (select number as tt from numbers(10)),  t2(tt) AS MATERIALIZED (S
 
 # https://github.com/databendlabs/databend/issues/17295
 statement ok
-create table t1(a int);
+create or replace table t1(a int);
 
 statement ok
-create table t2(a int);
+create or replace table t2(a int);
 
 query I
 WITH RECURSIVE
@@ -591,3 +591,32 @@ UNION all SELECT a from t2, closure where t2.a = closure.x
 ) SELECT x FROM closure;
 ----
 1
+
+# https://github.com/databendlabs/databend/issues/17432
+statement ok
+CREATE OR REPLACE TABLE some_table (
+	id INTEGER NOT NULL,
+	"data" VARCHAR(50),
+	parent_id INTEGER
+);
+
+statement ok
+INSERT INTO some_table (id, "data", parent_id) VALUES  (1, 'd1', NULL),(2, 'd2', 1),(3, 'd3', 1),(4, 'd4', 3),(5, 'd5', 3);
+
+query T
+WITH RECURSIVE some_cte(id, "data", parent_id) AS
+(SELECT some_table.id AS id, some_table."data" AS "data", some_table.parent_id AS parent_id
+FROM some_table
+WHERE some_table."data" IN ('d2', 'd3', 'd4') UNION ALL SELECT some_table_1.id AS id, some_table_1."data" AS "data", some_table_1.parent_id AS parent_id
+FROM some_table AS some_table_1, some_cte AS c1
+WHERE some_table_1.id = c1.parent_id)
+ SELECT some_cte."data"
+FROM some_cte
+WHERE some_cte."data" != 'd2';
+----
+d1
+d1
+d1
+d3
+d3
+d4