Skip to content

Commit 9986525

Browse files
yonatan-sevenaiclaude
yonatan-sevenai
and
claude
authored
fix(unparser): fold Limit/Sort into outer SELECT when Projection claims Aggregate through them (apache#21375)
## Which issue does this PR close? - Closes apache#21374 ## Rationale for this change When the SQL unparser encounters a Projection → Limit → Aggregate or Projection → Sort → Aggregate plan shape where aggregate aliases are inlined (no intermediate Projection between Limit/Sort and Aggregate), it emits the aggregate expressions twice - once in the outer SELECT and once inside a spurious derived subquery. The outer SELECT also references columns that are out of scope. This plan shape doesn't occur from the SQL parser (which inserts an intermediate Projection), but optimizers and plan builders can produce it. ## What changes are included in this PR? datafusion/sql/src/unparser/plan.rs: - reconstruct_select_statement now returns Result<bool> - true when it found and claimed an Aggregate node for the current SELECT. - In the Projection arm of select_to_sql_recursively, after claiming an Aggregate, if the Projection's direct child is a Limit or Sort, fold its clauses (LIMIT/OFFSET or ORDER BY/LIMIT) into the current query and recurse into the child's input, skipping the Limit/Sort node. This prevents the Limit/Sort arm from seeing already_projected and wrapping everything in a spurious derived subquery. datafusion/sql/tests/cases/plan_to_sql.rs: - roundtrip_aggregate_over_subquery — roundtrip test for the parser-generated plan shape (Projection between Limit and Aggregate), confirming it still works. - roundtrip_subquery_aggregate_with_column_alias - roundtrip test for SELECT id FROM (SELECT max(j1_id) FROM j1) AS c(id). - test_unparse_aggregate_over_subquery_no_inner_proj - manually constructed Projection → Limit → Aggregate plan, verifying the Limit is folded into the outer SELECT. - test_unparse_aggregate_no_outer_rename - manually constructed Projection → Aggregate with no outer rename, verifying aggregate aliases are preserved. - test_unparse_aggregate_with_sort_no_inner_proj - manually constructed Projection → Sort → Aggregate plan, verifying the Sort is folded into the outer SELECT. ## Are these changes tested? Yes. Five new tests covering: 1. Parser-generated roundtrip (regression guard) 2. Column alias roundtrip with subquery aggregate 3. Limit folding with manually constructed plan (was the primary bug) 4. Aggregate alias preservation without outer rename 5. Sort folding with manually constructed plan (same bug pattern) ## Are there any user-facing changes? No API changes. The fix corrects SQL output for programmatically constructed logical plans that have Projection → Limit → Aggregate or Projection → Sort → Aggregate shapes without an intermediate Projection. --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 05ea11e commit 9986525

2 files changed

Lines changed: 593 additions & 15 deletions

File tree

‎datafusion/sql/src/unparser/plan.rs‎

Lines changed: 267 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -45,12 +45,13 @@ use datafusion_common::{
4545
Column, DataFusionError, Result, ScalarValue, TableReference, assert_or_internal_err,
4646
internal_datafusion_err, internal_err, not_impl_err,
4747
tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRecursion},
48+
utils::combine_limit,
4849
};
4950
use datafusion_expr::expr::{OUTER_REFERENCE_COLUMN_PREFIX, UNNEST_COLUMN_PREFIX};
5051
use datafusion_expr::{
51-
Aggregate, BinaryExpr, Distinct, Expr, JoinConstraint, JoinType, LogicalPlan,
52-
LogicalPlanBuilder, Operator, Projection, SortExpr, TableScan, Unnest,
53-
UserDefinedLogicalNode, Window, expr::Alias,
52+
Aggregate, BinaryExpr, Distinct, Expr, FetchType, JoinConstraint, JoinType,
53+
LogicalPlan, LogicalPlanBuilder, Operator, Projection, SkipType, Sort, SortExpr,
54+
TableScan, Unnest, UserDefinedLogicalNode, Window, expr::Alias,
5455
};
5556
use sqlparser::ast::{self, Ident, OrderByKind, SetExpr, TableAliasColumnDef};
5657
use std::{sync::Arc, vec};
@@ -226,15 +227,29 @@ impl Unparser<'_> {
226227
Ok(SetExpr::Select(Box::new(select_builder.build()?)))
227228
}
228229

229-
/// Reconstructs a SELECT SQL statement from a logical plan by unprojecting column expressions
230-
/// found in a [Projection] node. This requires scanning the plan tree for relevant Aggregate
231-
/// and Window nodes and matching column expressions to the appropriate agg or window expressions.
230+
/// Reconstructs a SELECT SQL statement from a logical plan by
231+
/// unprojecting column expressions found in a [Projection] node. This
232+
/// requires scanning the plan tree for relevant Aggregate and Window
233+
/// nodes and matching column expressions to the appropriate agg or
234+
/// window expressions.
235+
///
236+
/// `fully_absorbed` reports whether the Projection arm was able to
237+
/// absorb every `Sort`/`Limit` node between this Projection and the
238+
/// Aggregate/Window into the current SELECT. When `false`, the
239+
/// Aggregate/Window will end up in a derived subquery, so we fall
240+
/// back to passthrough column references that resolve against that
241+
/// subquery's output instead of unprojecting onto the original
242+
/// aggregate expressions.
243+
///
244+
/// Returns `true` if an Aggregate node was found and claimed for this
245+
/// SELECT.
232246
fn reconstruct_select_statement(
233247
&self,
234248
plan: &LogicalPlan,
235249
p: &Projection,
236250
select: &mut SelectBuilder,
237-
) -> Result<()> {
251+
fully_absorbed: bool,
252+
) -> Result<bool> {
238253
let mut exprs = p.expr.clone();
239254

240255
// If an Unnest node is found within the select, find and unproject the unnest column
@@ -277,10 +292,24 @@ impl Unparser<'_> {
277292
.collect::<Result<Vec<_>>>()?;
278293
}
279294

280-
match (
281-
find_agg_node_within_select(plan, true),
282-
find_window_nodes_within_select(plan, None, true),
283-
) {
295+
// When some Sort/Limit nodes between this Projection and the
296+
// Aggregate/Window couldn't be absorbed into the current SELECT,
297+
// the Aggregate/Window will live inside a derived subquery. In
298+
// that case we use the passthrough projection path — column refs
299+
// resolve against the derived subquery's output columns instead
300+
// of being unprojected onto the original aggregate/window
301+
// expressions.
302+
let agg = if fully_absorbed {
303+
find_agg_node_within_select(plan, true)
304+
} else {
305+
None
306+
};
307+
let window = if fully_absorbed {
308+
find_window_nodes_within_select(plan, None, true)
309+
} else {
310+
None
311+
};
312+
match (agg, window) {
284313
(Some(agg), window) => {
285314
let window_option = window.as_deref();
286315
let items = exprs
@@ -299,6 +328,7 @@ impl Unparser<'_> {
299328
.collect::<Result<Vec<_>>>()?,
300329
vec![],
301330
));
331+
Ok(true)
302332
}
303333
(None, Some(window)) => {
304334
let items = exprs
@@ -310,6 +340,7 @@ impl Unparser<'_> {
310340
.collect::<Result<Vec<_>>>()?;
311341

312342
select.projection(items);
343+
Ok(false)
313344
}
314345
_ => {
315346
let items = exprs
@@ -328,9 +359,9 @@ impl Unparser<'_> {
328359
})
329360
.collect::<Result<Vec<_>>>()?;
330361
select.projection(items);
362+
Ok(false)
331363
}
332364
}
333-
Ok(())
334365
}
335366

336367
fn derive(
@@ -438,7 +469,7 @@ impl Unparser<'_> {
438469
}));
439470

440471
if !select.already_projected() {
441-
self.reconstruct_select_statement(plan, p, select)?;
472+
self.reconstruct_select_statement(plan, p, select, true)?;
442473
}
443474

444475
if matches!(
@@ -680,8 +711,229 @@ impl Unparser<'_> {
680711
if self.dialect.unnest_as_lateral_flatten() {
681712
Self::collect_flatten_aliases(p.input.as_ref(), select);
682713
}
683-
self.reconstruct_select_statement(plan, p, select)?;
684-
self.select_to_sql_recursively(p.input.as_ref(), query, select, relation)
714+
// Walk down through consecutive Sort/Limit nodes, greedily
715+
// absorbing what can be folded into the SELECT we're
716+
// building around the Aggregate. A single SQL SELECT can
717+
// carry at most one `ORDER BY` (applied before `LIMIT`),
718+
// so the safe shape between us and the Aggregate is
719+
// `Limit* Sort?` (outer→inner). We stop at the first node
720+
// that would violate this; that node becomes the
721+
// subquery boundary, and recursion (seeing
722+
// `already_projected = true`) wraps it in a derived
723+
// relation. If we walk all the way to a non-Sort/non-Limit
724+
// terminator, the entire chain folds into one SELECT.
725+
//
726+
// Stacked Sorts with nothing between them collapse to the
727+
// outermost — the same simplification `EnforceSorting`
728+
// applies on the physical side — but only when no Limit
729+
// has been absorbed since the previous Sort, since the
730+
// inner Sort would otherwise be determining which rows
731+
// the Limit keeps.
732+
//
733+
// The fold is collected here without touching `query`
734+
// (apart from non-literal direct Limits, which don't
735+
// depend on projection form). Once we know whether every
736+
// Sort/Limit was absorbed we can pick the right
737+
// projection form and emit `ORDER BY` with or without
738+
// unprojection.
739+
let mut cur = p.input.as_ref();
740+
let mut absorbed_sort: Option<&Sort> = None;
741+
let mut combined_skip: usize = 0;
742+
let mut combined_fetch: Option<usize> = None;
743+
let mut have_combined_limit = false;
744+
let mut have_direct_limit = false;
745+
let mut have_order_by = false;
746+
loop {
747+
match cur {
748+
LogicalPlan::Limit(limit) => {
749+
if have_order_by {
750+
// Limit-below-Sort: `ORDER BY … LIMIT N`
751+
// would apply the sort first, but the
752+
// logical plan applies the Limit first.
753+
break;
754+
}
755+
let skip_lit = limit.get_skip_type()?;
756+
let fetch_lit = limit.get_fetch_type()?;
757+
match (skip_lit, fetch_lit) {
758+
(SkipType::Literal(s), FetchType::Literal(f)) => {
759+
if have_direct_limit {
760+
break;
761+
}
762+
if have_combined_limit {
763+
// outer = already-accumulated;
764+
// inner = this Limit. Same merge
765+
// rule as the optimizer.
766+
let (cs, cf) = combine_limit(
767+
combined_skip,
768+
combined_fetch,
769+
s,
770+
f,
771+
);
772+
combined_skip = cs;
773+
combined_fetch = cf;
774+
} else {
775+
combined_skip = s;
776+
combined_fetch = f;
777+
have_combined_limit = true;
778+
}
779+
}
780+
_ => {
781+
if have_combined_limit || have_direct_limit {
782+
// Cannot safely merge a
783+
// non-literal Limit with a prior
784+
// one; let recursion handle it.
785+
break;
786+
}
787+
let Some(query_ref) = query.as_mut() else {
788+
return internal_err!(
789+
"Limit operator only valid in a statement context."
790+
);
791+
};
792+
if let Some(fetch) = &limit.fetch {
793+
query_ref.limit(Some(self.expr_to_sql(fetch)?));
794+
}
795+
if let Some(skip) = &limit.skip {
796+
query_ref.offset(Some(ast::Offset {
797+
rows: ast::OffsetRows::None,
798+
value: self.expr_to_sql(skip)?,
799+
}));
800+
}
801+
have_direct_limit = true;
802+
}
803+
}
804+
cur = limit.input.as_ref();
805+
}
806+
LogicalPlan::Sort(sort) if sort.fetch.is_some() => {
807+
// `Sort { fetch }` is logically
808+
// `Limit(fetch) -> Sort`. Try to absorb the
809+
// virtual Limit first; only if that succeeds
810+
// do we absorb the Sort. Otherwise we'd
811+
// silently drop the fetch.
812+
let fetch = sort.fetch.expect("guarded above");
813+
if have_order_by {
814+
// The virtual Limit would sit below an
815+
// already-absorbed outer Sort.
816+
break;
817+
}
818+
if have_direct_limit {
819+
// Cannot combine a literal fetch with a
820+
// non-literal direct Limit; let the
821+
// derived subquery preserve both.
822+
break;
823+
}
824+
if have_combined_limit {
825+
let (cs, cf) = combine_limit(
826+
combined_skip,
827+
combined_fetch,
828+
0,
829+
Some(fetch),
830+
);
831+
combined_skip = cs;
832+
combined_fetch = cf;
833+
} else {
834+
combined_skip = 0;
835+
combined_fetch = Some(fetch);
836+
have_combined_limit = true;
837+
}
838+
// Now the Sort itself. We know
839+
// `!have_order_by` from the check above.
840+
absorbed_sort = Some(sort);
841+
have_order_by = true;
842+
cur = sort.input.as_ref();
843+
}
844+
LogicalPlan::Sort(sort) => {
845+
// Sort without `fetch`.
846+
if have_order_by {
847+
// Outer Sort already absorbed; the inner
848+
// Sort is reordered by it and is
849+
// conventionally dropped, matching
850+
// `EnforceSorting` on the physical side.
851+
cur = sort.input.as_ref();
852+
continue;
853+
}
854+
absorbed_sort = Some(sort);
855+
have_order_by = true;
856+
cur = sort.input.as_ref();
857+
}
858+
_ => break,
859+
}
860+
}
861+
862+
// `fully_absorbed` is the bottom-up algorithm's "walked
863+
// all the way to the terminator without stopping": the
864+
// Aggregate/Window will live in the same SELECT as this
865+
// Projection, so we can unproject sort exprs and let
866+
// `reconstruct_select_statement` claim it.
867+
let fully_absorbed =
868+
!matches!(cur, LogicalPlan::Limit(_) | LogicalPlan::Sort(_));
869+
let found_agg =
870+
self.reconstruct_select_statement(plan, p, select, fully_absorbed)?;
871+
872+
// Whether to bother emitting the absorbed clauses: only
873+
// if there's an Aggregate either claimed in this SELECT
874+
// or about to live in a derived subquery below us. If
875+
// there's nothing aggregate-like to fold over, fall
876+
// through and let the normal recursion handle the
877+
// Projection's input.
878+
let agg_below =
879+
!fully_absorbed && find_agg_node_within_select(plan, true).is_some();
880+
if !(found_agg || agg_below) {
881+
return self.select_to_sql_recursively(
882+
p.input.as_ref(),
883+
query,
884+
select,
885+
relation,
886+
);
887+
}
888+
889+
if let Some(sort) = absorbed_sort {
890+
let Some(query_ref) = query.as_mut() else {
891+
return internal_err!(
892+
"Sort operator only valid in a statement context."
893+
);
894+
};
895+
let sort_exprs: Vec<SortExpr> = if fully_absorbed {
896+
let agg =
897+
find_agg_node_within_select(plan, select.already_projected());
898+
sort.expr
899+
.iter()
900+
.map(|sort_expr| {
901+
unproject_sort_expr(
902+
sort_expr.clone(),
903+
agg,
904+
sort.input.as_ref(),
905+
)
906+
})
907+
.collect::<Result<Vec<_>>>()?
908+
} else {
909+
sort.expr.clone()
910+
};
911+
query_ref.order_by(self.sorts_to_sql(&sort_exprs)?);
912+
}
913+
if have_combined_limit {
914+
let Some(query_ref) = query.as_mut() else {
915+
return internal_err!(
916+
"Limit operator only valid in a statement context."
917+
);
918+
};
919+
if let Some(fetch) = combined_fetch {
920+
query_ref.limit(Some(ast::Expr::value(ast::Value::Number(
921+
fetch.to_string(),
922+
false,
923+
))));
924+
}
925+
if combined_skip > 0 {
926+
query_ref.offset(Some(ast::Offset {
927+
rows: ast::OffsetRows::None,
928+
value: ast::Expr::value(ast::Value::Number(
929+
combined_skip.to_string(),
930+
false,
931+
)),
932+
}));
933+
}
934+
}
935+
936+
self.select_to_sql_recursively(cur, query, select, relation)
685937
}
686938
LogicalPlan::Filter(filter) => {
687939
let window = find_window_nodes_within_select(

0 commit comments

Comments
 (0)