Skip to content

Commit 4dadbbd

Browse files
authored
fix: apply recursive CTE column-list aliases to the static term (apache#23098)
## Which issue does this PR close? - Closes apache#23097. ## Rationale for this change `WITH RECURSIVE t(n) AS (...)` failed to plan because the CTE's declared column-list names (the `t(n)` part) were never applied to the recursive working relation. They were applied (via `apply_table_alias`) only after the whole CTE plan was built, but the working table is derived from the static term's schema *before* that — so the self-reference couldn't resolve the declared names and planning failed with `Schema error: No field named n. Valid fields are t."Int64(1)".`. PostgreSQL and DuckDB accept the query; aliasing inside the static `SELECT` (`SELECT 1 AS n`) was the only workaround. ## What changes are included in this PR? Apply the column-list aliases to the static term inside `recursive_cte()`, before the work table is created, so the working relation and the self-reference carry the declared names. The caller now applies only the relation-name alias on the recursive path (the columns are already applied), avoiding a redundant projection on top of the `RecursiveQuery` node. The non-UNION fallback applies the aliases directly; non-recursive CTEs are unchanged. A column/alias-count mismatch is now reported at the static term — a clearer error than the previous "No field named …". ## Are these changes tested? Yes, added `cte.slt` cases for single- and multi-column column-list recursive CTEs (asserting the recursion produces the expected rows), `UNION (DISTINCT)`, the arity-mismatch error, and an `EXPLAIN` locking the plan shape (no extra projection over `RecursiveQuery`). ## Are there any user-facing changes? `WITH RECURSIVE t(n) AS (...)` and multi-column column lists now plan and execute correctly, matching PostgreSQL/DuckDB. No API changes. *Note: this pull request was created together with AI tools (claude code), the full diff was reviewed by myself in full prior to submission*
1 parent 6ef3a3b commit 4dadbbd

3 files changed

Lines changed: 140 additions & 18 deletions

File tree

datafusion/optimizer/tests/optimizer_integration.rs

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -74,20 +74,20 @@ fn recursive_cte_with_nested_subquery() -> Result<()> {
7474

7575
assert_snapshot!(
7676
format!("{plan}"),
77-
@r"
77+
@"
7878
SubqueryAlias: numbers
79-
Projection: sub.id AS id, sub.level AS level
80-
RecursiveQuery: is_distinct=false
79+
RecursiveQuery: is_distinct=false
80+
Projection: sub.id AS id, sub.level AS level
8181
SubqueryAlias: sub
8282
Projection: test.col_int32 AS id, Int64(1) AS level
8383
TableScan: test projection=[col_int32]
84-
Projection: t.col_int32, numbers.level + Int64(1)
85-
Inner Join: CAST(t.col_int32 AS Int64) = CAST(numbers.id AS Int64) + Int64(1)
86-
SubqueryAlias: t
87-
Filter: CAST(test.col_int32 AS Int64) IS NOT NULL
88-
TableScan: test projection=[col_int32]
89-
Filter: CAST(numbers.id AS Int64) + Int64(1) IS NOT NULL
90-
TableScan: numbers projection=[id, level]
84+
Projection: t.col_int32, numbers.level + Int64(1)
85+
Inner Join: CAST(t.col_int32 AS Int64) = CAST(numbers.id AS Int64) + Int64(1)
86+
SubqueryAlias: t
87+
Filter: CAST(test.col_int32 AS Int64) IS NOT NULL
88+
TableScan: test projection=[col_int32]
89+
Filter: CAST(numbers.id AS Int64) + Int64(1) IS NOT NULL
90+
TableScan: numbers projection=[id, level]
9191
"
9292
);
9393

datafusion/sql/src/cte.rs

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,11 @@ use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
2121

2222
use arrow::datatypes::{Schema, SchemaRef};
2323
use datafusion_common::{
24-
Result, not_impl_err, plan_err,
24+
Result, TableReference, not_impl_err, plan_err,
2525
tree_node::{TreeNode, TreeNodeRecursion},
2626
};
2727
use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, TableSource};
28-
use sqlparser::ast::{Query, SetExpr, SetOperator, With};
28+
use sqlparser::ast::{Ident, Query, SetExpr, SetOperator, With};
2929

3030
impl<S: ContextProvider> SqlToRel<'_, S> {
3131
pub(super) fn plan_with_clause(
@@ -46,14 +46,24 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
4646

4747
// Create a logical plan for the CTE
4848
let cte_plan = if is_recursive {
49-
self.recursive_cte(&cte_name, *cte.query, planner_context)?
49+
let columns = cte.alias.columns.iter().map(|c| c.name.clone()).collect();
50+
self.recursive_cte(&cte_name, columns, *cte.query, planner_context)?
5051
} else {
5152
self.non_recursive_cte(*cte.query, planner_context)?
5253
};
5354

54-
// Each `WITH` block can change the column names in the last
55-
// projection (e.g. "WITH table(t1, t2) AS SELECT 1, 2").
56-
let final_plan = self.apply_table_alias(cte_plan, cte.alias)?;
55+
// Each `WITH` block can change the column names in the last projection
56+
// (e.g. "WITH table(t1, t2) AS SELECT 1, 2"). Recursive CTEs apply those
57+
// to the static term in recursive_cte(), so only the relation name here.
58+
let final_plan = if is_recursive {
59+
LogicalPlanBuilder::from(cte_plan)
60+
.alias(TableReference::bare(
61+
self.ident_normalizer.normalize(cte.alias.name),
62+
))?
63+
.build()?
64+
} else {
65+
self.apply_table_alias(cte_plan, cte.alias)?
66+
};
5767
// Export the CTE to the outer query
5868
planner_context.insert_cte(cte_name, final_plan);
5969
}
@@ -71,6 +81,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
7181
fn recursive_cte(
7282
&self,
7383
cte_name: &str,
84+
columns: Vec<Ident>,
7485
mut cte_query: Query,
7586
planner_context: &mut PlannerContext,
7687
) -> Result<LogicalPlan> {
@@ -91,9 +102,11 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
91102
set_quantifier,
92103
} => (left, right, set_quantifier),
93104
other => {
94-
// If the query is not a UNION, then it is not a recursive CTE
105+
// Not a UNION, so not actually a recursive CTE. The caller adds only
106+
// the relation name for recursive CTEs, so apply the column aliases here.
95107
*cte_query.body = other;
96-
return self.non_recursive_cte(cte_query, planner_context);
108+
let plan = self.non_recursive_cte(cte_query, planner_context)?;
109+
return self.apply_expr_alias(plan, columns);
97110
}
98111
};
99112

@@ -111,6 +124,10 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
111124
// ---------- Step 1: Compile the static term ------------------
112125
let static_plan = self.set_expr_to_plan(*left_expr, planner_context)?;
113126

127+
// Apply the declared column-list aliases (e.g. `t(n)`) to the static term, so
128+
// the work table built from its schema below exposes the declared names.
129+
let static_plan = self.apply_expr_alias(static_plan, columns)?;
130+
114131
// Since the recursive CTEs include a component that references a
115132
// table with its name, like the example below:
116133
//

datafusion/sqllogictest/test_files/cte.slt

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,111 @@ physical_plan
179179
07)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
180180
08)----------WorkTableExec: name=nodes
181181

182+
# recursive CTE with a column-list alias (e.g. `t(n)`): the declared names must be
183+
# applied to the static term so the recursive self-reference can resolve them
184+
query I rowsort
185+
WITH RECURSIVE t(n) AS (
186+
SELECT 1
187+
UNION ALL
188+
SELECT n + 1 FROM t WHERE n < 10
189+
)
190+
SELECT n FROM t
191+
----
192+
1
193+
10
194+
2
195+
3
196+
4
197+
5
198+
6
199+
7
200+
8
201+
9
202+
203+
# recursive CTE with a multi-column column-list alias
204+
query II rowsort
205+
WITH RECURSIVE t(a, b) AS (
206+
SELECT 1, 2
207+
UNION ALL
208+
SELECT a + 1, b * 2 FROM t WHERE a < 5
209+
)
210+
SELECT a, b FROM t
211+
----
212+
1 2
213+
2 4
214+
3 8
215+
4 16
216+
5 32
217+
218+
# recursive CTE with a column-list alias and UNION (DISTINCT)
219+
query I rowsort
220+
WITH RECURSIVE t(n) AS (
221+
SELECT 1
222+
UNION
223+
SELECT n + 1 FROM t WHERE n < 5
224+
)
225+
SELECT n FROM t
226+
----
227+
1
228+
2
229+
3
230+
4
231+
5
232+
233+
# recursive CTE column-list alias arity mismatch is rejected cleanly (raised at
234+
# the static term, rather than the old confusing "No field named ...")
235+
query error DataFusion error: Error during planning: Source table contains 1 columns but only 2 names given as column alias
236+
WITH RECURSIVE t(a, b) AS (
237+
SELECT 1
238+
UNION ALL
239+
SELECT a + 1 FROM t WHERE a < 3
240+
)
241+
SELECT * FROM t
242+
243+
# explain a column-list-aliased recursive CTE: the declared name is applied to
244+
# the static term, so there is no extra projection on top of RecursiveQuery
245+
query TT
246+
EXPLAIN WITH RECURSIVE t(n) AS (
247+
SELECT 1
248+
UNION ALL
249+
SELECT n + 1 FROM t WHERE n < 10
250+
)
251+
SELECT * FROM t
252+
----
253+
logical_plan
254+
01)SubqueryAlias: t
255+
02)--RecursiveQuery: is_distinct=false
256+
03)----Projection: Int64(1) AS n
257+
04)------EmptyRelation: rows=1
258+
05)----Projection: t.n + Int64(1)
259+
06)------Filter: t.n < Int64(10)
260+
07)--------TableScan: t projection=[n]
261+
physical_plan
262+
01)RecursiveQueryExec: name=t, is_distinct=false
263+
02)--ProjectionExec: expr=[CAST(1 AS Int64) as n]
264+
03)----PlaceholderRowExec
265+
04)--CoalescePartitionsExec
266+
05)----ProjectionExec: expr=[n@0 + 1 as n]
267+
06)------FilterExec: n@0 < 10
268+
07)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
269+
08)----------WorkTableExec: name=t
270+
271+
# recursive CTE with a quoted, case-sensitive column-list alias: `"N"` must be
272+
# preserved (not lowercased) so the recursive self-reference resolves it
273+
query I rowsort
274+
WITH RECURSIVE t("N") AS (
275+
SELECT 1
276+
UNION ALL
277+
SELECT "N" + 1 FROM t WHERE "N" < 5
278+
)
279+
SELECT "N" FROM t
280+
----
281+
1
282+
2
283+
3
284+
4
285+
5
286+
182287
# simple deduplicating recursive CTE works
183288
query I
184289
WITH RECURSIVE nodes AS (

0 commit comments

Comments
 (0)