diff --git a/src/query/sql/src/executor/physical_plans/physical_union_all.rs b/src/query/sql/src/executor/physical_plans/physical_union_all.rs index 41caccea8bb9a..1e7c07829e8bf 100644 --- a/src/query/sql/src/executor/physical_plans/physical_union_all.rs +++ b/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 = (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 = (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. diff --git a/src/query/sql/src/planner/binder/util.rs b/src/query/sql/src/planner/binder/util.rs index cd99be2f3f193..82e9981bcb5c9 100644 --- a/src/query/sql/src/planner/binder/util.rs +++ b/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::>(), - ); + cte_types.extend(plan.fields.iter().map(|f| f.data_type().clone())); } } diff --git a/tests/sqllogictests/suites/query/cte/cte.test b/tests/sqllogictests/suites/query/cte/cte.test index a589ccb007c00..cd70e3935184a 100644 --- a/tests/sqllogictests/suites/query/cte/cte.test +++ b/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