From 5483a1d3e189b739b2ced924babc81c4974c81e2 Mon Sep 17 00:00:00 2001 From: naman-modi Date: Wed, 22 Jul 2026 11:17:38 +0530 Subject: [PATCH 1/2] refactor(unparser): centralize aggregate-scope rendering in the SQL unparser --- datafusion/core/tests/sql/unparser.rs | 133 +++++++++++++++++++++++ datafusion/sql/src/unparser/plan.rs | 150 ++++++++++++++++---------- 2 files changed, 227 insertions(+), 56 deletions(-) diff --git a/datafusion/core/tests/sql/unparser.rs b/datafusion/core/tests/sql/unparser.rs index d689fb1496a2b..75ecaa4f310f9 100644 --- a/datafusion/core/tests/sql/unparser.rs +++ b/datafusion/core/tests/sql/unparser.rs @@ -468,6 +468,58 @@ QUALIFY rn = 1 AND count(DISTINCT cs.customer_id) > 0 "#; +// https://github.com/apache/datafusion/issues/23668 +// +// Extends the #23317 aggregate-scope fix to the window and ORDER BY clauses. +// Reuses issue_23317_context() (same derived-projection shape). + +// Window sorting by an aggregate, over a derived-projection input. Already +// correct today; this locks the OVER clause against keeping the out-of-scope +// `cs` qualifier across the refactor. +const ISSUE_23668_WINDOW_QUERY: &str = r#" +SELECT + date_part('year', c.signup_date) AS signup_year, + count(DISTINCT cs.customer_id) AS customers, + row_number() OVER (ORDER BY count(DISTINCT cs.customer_id) DESC) AS rn +FROM + "warehouse"."main"."sales" cs + JOIN "warehouse"."main"."customers" c USING (customer_id) +GROUP BY + 1 +"#; + +// ORDER BY an aggregate that is NOT selected, so it can't use a select alias +// and is unprojected through the Aggregate. It must be normalized like the +// SELECT list, not keep the out-of-scope `cs` qualifier. +const ISSUE_23668_ORDER_BY_QUERY: &str = r#" +SELECT + date_part('year', c.signup_date) AS signup_year, + count(DISTINCT cs.customer_id) AS customers +FROM + "warehouse"."main"."sales" cs + JOIN "warehouse"."main"."customers" c USING (customer_id) +GROUP BY + 1 +ORDER BY + round(sum(cs.total_revenue), 2) DESC +"#; + +// ORDER BY a selected aggregate: the optimizer keeps a top-level Sort, so this +// covers the direct `Sort` arm (the other sort caller, vs the projection- +// absorbed one above). +const ISSUE_23668_TOP_LEVEL_SORT_QUERY: &str = r#" +SELECT + date_part('year', c.signup_date) AS signup_year, + count(DISTINCT cs.customer_id) AS customers +FROM + "warehouse"."main"."sales" cs + JOIN "warehouse"."main"."customers" c USING (customer_id) +GROUP BY + 1 +ORDER BY + customers DESC +"#; + fn issue_23317_context() -> Result { let ctx = SessionContext::new(); @@ -612,6 +664,87 @@ async fn optimized_duckdb_unparse_qualify_unqualifies_agg_input() -> Result<()> Ok(()) } +#[tokio::test] +async fn optimized_duckdb_unparse_window_over_agg_unqualifies_input() -> Result<()> { + let ctx = issue_23317_context()?; + assert!(ctx.remove_optimizer_rule(SingleDistinctToGroupBy::new().name())); + + let plan = ctx + .sql(ISSUE_23668_WINDOW_QUERY) + .await? + .into_optimized_plan()?; + let dialect = DuckDBDialect::new(); + let unparser = Unparser::new(&dialect); + let sql = unparser.plan_to_sql(&plan)?.to_string(); + + assert_issue_23317_unparsed_sql_plans(&ctx, &sql).await?; + + assert!( + sql.contains(r#"OVER (ORDER BY count(DISTINCT "customer_id")"#), + "window ORDER BY aggregate should resolve against the derived projection output: {sql}", + ); + assert!( + !sql.contains(r#"count(DISTINCT "cs"."customer_id")"#), + "window OVER clause must not reference out-of-scope alias cs: {sql}", + ); + + Ok(()) +} + +#[tokio::test] +async fn optimized_duckdb_unparse_order_by_unqualifies_agg_input() -> Result<()> { + let ctx = issue_23317_context()?; + assert!(ctx.remove_optimizer_rule(SingleDistinctToGroupBy::new().name())); + + let plan = ctx + .sql(ISSUE_23668_ORDER_BY_QUERY) + .await? + .into_optimized_plan()?; + let dialect = DuckDBDialect::new(); + let unparser = Unparser::new(&dialect); + let sql = unparser.plan_to_sql(&plan)?.to_string(); + + assert_issue_23317_unparsed_sql_plans(&ctx, &sql).await?; + + assert!( + sql.contains(r#"ORDER BY round(sum("total_revenue"), 2)"#), + "ORDER BY aggregate should resolve against the derived projection output: {sql}", + ); + assert!( + !sql.contains(r#"sum("cs"."total_revenue")"#), + "ORDER BY must not reference out-of-scope alias cs: {sql}", + ); + + Ok(()) +} + +#[tokio::test] +async fn optimized_duckdb_unparse_top_level_sort_over_agg_stays_in_scope() -> Result<()> { + let ctx = issue_23317_context()?; + assert!(ctx.remove_optimizer_rule(SingleDistinctToGroupBy::new().name())); + + let plan = ctx + .sql(ISSUE_23668_TOP_LEVEL_SORT_QUERY) + .await? + .into_optimized_plan()?; + let dialect = DuckDBDialect::new(); + let unparser = Unparser::new(&dialect); + let sql = unparser.plan_to_sql(&plan)?.to_string(); + + assert_issue_23317_unparsed_sql_plans(&ctx, &sql).await?; + + assert!( + sql.contains(r#"ORDER BY "customers""#), + "top-level ORDER BY should resolve to the select alias: {sql}", + ); + assert!( + !sql.contains(r#""cs"."customer_id") AS "customers""#), + "aggregate output must not reference out-of-scope alias cs: {sql}", + ); + + Ok(()) +} + /// The outcome of running a single roundtrip test. /// /// A successful test produces [`TestCaseResult::Success`]. diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 5eef9b82d975e..719b9510f5aff 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -101,6 +101,75 @@ pub fn plan_to_sql(plan: &LogicalPlan) -> Result { unparser.plan_to_sql(plan) } +/// Aggregate-expression scope for one rendered SELECT block. +/// +/// When an aggregate's input is itself emitted as a derived subquery (a +/// projection sits between the aggregate and its relation), the input columns +/// are only reachable by that derived table's output names. Base-table +/// qualifiers like `t.col` name a relation that is out of scope above the +/// boundary, so emitting them produces SQL a strict engine rejects. +/// +/// Every clause that renders an aggregate expression (SELECT / GROUP BY / +/// HAVING / QUALIFY / ORDER BY) has to apply the same rule. Detect the +/// boundary once here and reuse it, so the clauses can't drift apart (which is +/// how earlier fixes left some clauses correct and others not). +struct UnparserAggScope<'a> { + agg: &'a Aggregate, + /// `agg.input` renders as a derived projection, so out-of-scope qualifiers + /// must be stripped from expressions in this scope. + input_is_derived_projection: bool, +} + +impl<'a> UnparserAggScope<'a> { + fn new(agg: &'a Aggregate) -> Self { + Self { + agg, + input_is_derived_projection: Unparser::contains_projection_before_relation( + agg.input.as_ref(), + ), + } + } + + /// Prepare a projected column or predicate that still references the + /// aggregate by its output columns: unproject it back onto the aggregate + /// (and `windows`) expressions, then normalize it for this scope. + fn prepare(&self, expr: Expr, windows: Option<&[&Window]>) -> Result { + self.normalize(unproject_agg_exprs(expr, self.agg, windows)?) + } + + /// Normalize an expression that is already in aggregate form (group / aggr + /// exprs, or an unprojected sort expr): strip the qualifiers that fall out + /// of scope once the input is a derived projection. No-op otherwise. + fn normalize(&self, expr: Expr) -> Result { + if self.input_is_derived_projection { + Unparser::strip_column_qualifiers_for_schema( + expr, + self.agg.input.schema().as_ref(), + ) + } else { + Ok(expr) + } + } + + /// Unproject a sort expression onto its aggregate, then normalize it so + /// ORDER BY resolves against the same scope as every other clause in the + /// block. Associated fn taking `Option<&Self>` because a sort can appear + /// with no aggregate below it, in which case only unprojection applies. + fn prepare_sort_expr( + scope: Option<&Self>, + sort_expr: SortExpr, + input: &LogicalPlan, + ) -> Result { + let mut sort_expr = unproject_sort_expr(sort_expr, scope.map(|s| s.agg), input)?; + + if let Some(scope) = scope { + sort_expr.expr = scope.normalize(sort_expr.expr)?; + } + + Ok(sort_expr) + } +} + impl Unparser<'_> { pub fn plan_to_sql(&self, plan: &LogicalPlan) -> Result { let mut plan = normalize_union_schema(plan)?; @@ -312,17 +381,12 @@ impl Unparser<'_> { match (agg, window) { (Some(agg), window) => { let window_option = window.as_deref(); - let agg_input_has_derived_projection = - Self::contains_projection_before_relation(agg.input.as_ref()); + let unparser_agg_scope = UnparserAggScope::new(agg); let items = exprs .into_iter() .map(|proj_expr| { - let unproj = unproject_agg_exprs(proj_expr, agg, window_option)?; - let unproj = Self::normalize_agg_input_columns( - unproj, - agg, - agg_input_has_derived_projection, - )?; + let unproj = + unparser_agg_scope.prepare(proj_expr, window_option)?; self.select_item_to_sql(&unproj) }) .collect::>>()?; @@ -333,12 +397,7 @@ impl Unparser<'_> { .iter() .cloned() .map(|expr| { - let expr = Self::normalize_agg_input_columns( - expr, - agg, - agg_input_has_derived_projection, - )?; - self.expr_to_sql(&expr) + self.expr_to_sql(&unparser_agg_scope.normalize(expr)?) }) .collect::>>()?, vec![], @@ -379,18 +438,6 @@ impl Unparser<'_> { } } - fn normalize_agg_input_columns( - expr: Expr, - agg: &Aggregate, - input_has_derived_projection: bool, - ) -> Result { - if input_has_derived_projection { - Self::strip_column_qualifiers_for_schema(expr, agg.input.schema().as_ref()) - } else { - Ok(expr) - } - } - fn contains_projection_before_relation(plan: &LogicalPlan) -> bool { match plan { LogicalPlan::Projection(_) => true, @@ -592,6 +639,11 @@ impl Unparser<'_> { window_expr .iter() .map(|expr| { + // No scope normalization here. This agg-carrying branch is + // only reached by a manually built plan (a Window with no + // enclosing Projection). SQL always wraps a window function + // in a projection, so real queries render through + // reconstruct_select_statement, which normalizes instead. let expr = if let Some(agg) = agg { unproject_agg_exprs(expr.clone(), agg, None)? } else { @@ -974,12 +1026,13 @@ impl Unparser<'_> { let sort_exprs: Vec = if fully_absorbed { let agg = find_agg_node_within_select(plan, select.already_projected()); + let unparser_agg_scope = agg.map(UnparserAggScope::new); sort.expr .iter() .map(|sort_expr| { - unproject_sort_expr( + UnparserAggScope::prepare_sort_expr( + unparser_agg_scope.as_ref(), sort_expr.clone(), - agg, sort.input.as_ref(), ) }) @@ -1028,23 +1081,14 @@ impl Unparser<'_> { let mut unprojected = unproject_window_exprs(filter.predicate.clone(), window)?; if let Some(agg) = agg { - unprojected = unproject_agg_exprs(unprojected, agg, None)?; - unprojected = Self::normalize_agg_input_columns( - unprojected, - agg, - Self::contains_projection_before_relation(agg.input.as_ref()), - )?; + unprojected = + UnparserAggScope::new(agg).prepare(unprojected, None)?; } let filter_expr = self.expr_to_sql(&unprojected)?; select.qualify(Some(filter_expr)); } else if let Some(agg) = agg { - let unprojected = - unproject_agg_exprs(filter.predicate.clone(), agg, None)?; - let unprojected = Self::normalize_agg_input_columns( - unprojected, - agg, - Self::contains_projection_before_relation(agg.input.as_ref()), - )?; + let unprojected = UnparserAggScope::new(agg) + .prepare(filter.predicate.clone(), None)?; let filter_expr = self.expr_to_sql(&unprojected)?; select.having(Some(filter_expr)); } else { @@ -1125,12 +1169,17 @@ impl Unparser<'_> { }; let agg = find_agg_node_within_select(plan, select.already_projected()); + let scope = agg.map(UnparserAggScope::new); // unproject sort expressions let sort_exprs: Vec = sort .expr .iter() .map(|sort_expr| { - unproject_sort_expr(sort_expr.clone(), agg, sort.input.as_ref()) + UnparserAggScope::prepare_sort_expr( + scope.as_ref(), + sort_expr.clone(), + sort.input.as_ref(), + ) }) .collect::>>()?; @@ -1146,8 +1195,7 @@ impl Unparser<'_> { LogicalPlan::Aggregate(agg) => { // Aggregation can be already handled in the projection case if !select.already_projected() { - let agg_input_has_derived_projection = - Self::contains_projection_before_relation(agg.input.as_ref()); + let unparser_agg_scope = UnparserAggScope::new(agg); // The query returns aggregate and group expressions. If that weren't the case, // the aggregate would have been placed inside a projection, making the check above^ false let exprs: Vec<_> = agg @@ -1156,12 +1204,7 @@ impl Unparser<'_> { .chain(agg.group_expr.iter()) .cloned() .map(|expr| { - let expr = Self::normalize_agg_input_columns( - expr, - agg, - agg_input_has_derived_projection, - )?; - self.select_item_to_sql(&expr) + self.select_item_to_sql(&unparser_agg_scope.normalize(expr)?) }) .collect::>>()?; select.projection(exprs); @@ -1171,12 +1214,7 @@ impl Unparser<'_> { .iter() .cloned() .map(|expr| { - let expr = Self::normalize_agg_input_columns( - expr, - agg, - agg_input_has_derived_projection, - )?; - self.expr_to_sql(&expr) + self.expr_to_sql(&unparser_agg_scope.normalize(expr)?) }) .collect::>>()?, vec![], From ef910d7d13870047b966de8fee36d1676492ce0b Mon Sep 17 00:00:00 2001 From: naman-modi Date: Thu, 23 Jul 2026 17:31:45 +0530 Subject: [PATCH 2/2] nit(refactor unparser): address review comments --- datafusion/core/tests/sql/unparser.rs | 10 +++--- datafusion/sql/src/unparser/plan.rs | 47 ++++++++++++++------------- 2 files changed, 31 insertions(+), 26 deletions(-) diff --git a/datafusion/core/tests/sql/unparser.rs b/datafusion/core/tests/sql/unparser.rs index 75ecaa4f310f9..355a58fd6f45b 100644 --- a/datafusion/core/tests/sql/unparser.rs +++ b/datafusion/core/tests/sql/unparser.rs @@ -504,9 +504,10 @@ ORDER BY round(sum(cs.total_revenue), 2) DESC "#; -// ORDER BY a selected aggregate: the optimizer keeps a top-level Sort, so this -// covers the direct `Sort` arm (the other sort caller, vs the projection- -// absorbed one above). +// ORDER BY a selected aggregate keeps a top-level Sort (the direct `Sort` arm, +// vs the projection-absorbed one above). It resolves to the select alias, so +// this covers routing only -- the normalization in that arm isn't reachable +// from SQL (an unselected aggregate takes the absorbed path above instead). const ISSUE_23668_TOP_LEVEL_SORT_QUERY: &str = r#" SELECT date_part('year', c.signup_date) AS signup_year, @@ -719,7 +720,8 @@ async fn optimized_duckdb_unparse_order_by_unqualifies_agg_input() -> Result<()> } #[tokio::test] -async fn optimized_duckdb_unparse_top_level_sort_over_agg_stays_in_scope() -> Result<()> { +async fn optimized_duckdb_unparse_top_level_sort_over_agg_uses_select_alias() -> Result<()> +{ let ctx = issue_23317_context()?; assert!(ctx.remove_optimizer_rule(SingleDistinctToGroupBy::new().name())); diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 719b9510f5aff..9fe97a8291b6a 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -151,21 +151,15 @@ impl<'a> UnparserAggScope<'a> { } } - /// Unproject a sort expression onto its aggregate, then normalize it so - /// ORDER BY resolves against the same scope as every other clause in the - /// block. Associated fn taking `Option<&Self>` because a sort can appear - /// with no aggregate below it, in which case only unprojection applies. + /// Unproject a sort expression onto this aggregate, then normalize it so + /// ORDER BY uses the same scope as the other clauses. fn prepare_sort_expr( - scope: Option<&Self>, + &self, sort_expr: SortExpr, input: &LogicalPlan, ) -> Result { - let mut sort_expr = unproject_sort_expr(sort_expr, scope.map(|s| s.agg), input)?; - - if let Some(scope) = scope { - sort_expr.expr = scope.normalize(sort_expr.expr)?; - } - + let mut sort_expr = unproject_sort_expr(sort_expr, Some(self.agg), input)?; + sort_expr.expr = self.normalize(sort_expr.expr)?; Ok(sort_expr) } } @@ -476,6 +470,19 @@ impl Unparser<'_> { } } + /// Unproject a sort expression; normalize it when the sort is above an + /// aggregate, otherwise just unproject (no scope to normalize against). + fn unproject_sort_expr_in_scope( + sort_expr: SortExpr, + agg: Option<&Aggregate>, + input: &LogicalPlan, + ) -> Result { + match agg { + Some(agg) => UnparserAggScope::new(agg).prepare_sort_expr(sort_expr, input), + None => unproject_sort_expr(sort_expr, None, input), + } + } + fn derive( &self, plan: &LogicalPlan, @@ -639,11 +646,9 @@ impl Unparser<'_> { window_expr .iter() .map(|expr| { - // No scope normalization here. This agg-carrying branch is - // only reached by a manually built plan (a Window with no - // enclosing Projection). SQL always wraps a window function - // in a projection, so real queries render through - // reconstruct_select_statement, which normalizes instead. + // No normalization: this agg branch is only reachable from a + // hand-built plan. SQL wraps windows in a projection, which + // reconstruct_select_statement handles (and normalizes). let expr = if let Some(agg) = agg { unproject_agg_exprs(expr.clone(), agg, None)? } else { @@ -1026,13 +1031,12 @@ impl Unparser<'_> { let sort_exprs: Vec = if fully_absorbed { let agg = find_agg_node_within_select(plan, select.already_projected()); - let unparser_agg_scope = agg.map(UnparserAggScope::new); sort.expr .iter() .map(|sort_expr| { - UnparserAggScope::prepare_sort_expr( - unparser_agg_scope.as_ref(), + Self::unproject_sort_expr_in_scope( sort_expr.clone(), + agg, sort.input.as_ref(), ) }) @@ -1169,15 +1173,14 @@ impl Unparser<'_> { }; let agg = find_agg_node_within_select(plan, select.already_projected()); - let scope = agg.map(UnparserAggScope::new); // unproject sort expressions let sort_exprs: Vec = sort .expr .iter() .map(|sort_expr| { - UnparserAggScope::prepare_sort_expr( - scope.as_ref(), + Self::unproject_sort_expr_in_scope( sort_expr.clone(), + agg, sort.input.as_ref(), ) })