From 947b9416188968dfb26e99a988649a6319f61f5f Mon Sep 17 00:00:00 2001 From: claudespice <270518434+claudespice@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:39:23 -0700 Subject: [PATCH] fix(unparser): unparse a stacked aggregate as a derived table MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A SELECT expresses a single grouping, but the `LogicalPlan::Aggregate` arm of `select_to_sql_recursively` recursed straight into its input whenever the select list was already built, so a second aggregate underneath was skipped and its `GROUP BY` never reached the emitted SQL. `single_distinct_to_groupby` produces exactly that shape for `count(DISTINCT c)`: an outer `count(alias1)` over an inner `Aggregate` grouping by `c AS alias1`. With the inner aggregate dropped, `SELECT count(DISTINCT "UserID") FROM hits` unparses to `SELECT count(alias1) FROM hits` — `alias1` does not exist on the base table, so a consumer pushing the optimized plan down to a remote engine gets a binder error, and where a column of that name does exist it counts every row rather than the distinct values. Track on the `SelectBuilder` whether an aggregate has been folded into the current SELECT, and emit a `derived_aggregate` table for the next one, as the `Sort`, `Limit`, `Distinct` and `Projection` arms already do. The existing TPC-H and ClickBench roundtrip suites unparse the plan as the SQL planner produces it, before `single_distinct_to_groupby` runs, which is why they never caught this; the new core test unparses the optimized plan and executes both statements. --- datafusion/core/tests/sql/unparser.rs | 37 ++++++++++++ datafusion/sql/src/unparser/ast.rs | 17 ++++++ datafusion/sql/src/unparser/plan.rs | 18 ++++++ datafusion/sql/tests/cases/plan_to_sql.rs | 69 +++++++++++++++++++++++ 4 files changed, 141 insertions(+) diff --git a/datafusion/core/tests/sql/unparser.rs b/datafusion/core/tests/sql/unparser.rs index d6ca872e198c3..3a208993a1b19 100644 --- a/datafusion/core/tests/sql/unparser.rs +++ b/datafusion/core/tests/sql/unparser.rs @@ -464,3 +464,40 @@ async fn test_tpch_unparser_roundtrip() { set_stack_allocation_size(8 * 1024 * 1024); run_roundtrip_tests("TPC-H", tpch_queries(), tpch_test_context).await; } + +/// The suites above unparse the plan as it comes out of the SQL planner. A consumer that +/// pushes a query down to a remote engine unparses the *optimized* plan, where +/// `single_distinct_to_groupby` has rewritten `count(DISTINCT c)` into an outer +/// `count(alias1)` over an inner `GROUP BY c AS alias1`. Both aggregates have to reach the +/// emitted SQL: with the inner one dropped, `count(alias1)` is left referencing the base +/// table, which fails to bind — and where a column of that name does exist, counts every +/// row instead of the distinct values. +#[tokio::test] +async fn test_optimized_plan_roundtrip_count_distinct() -> Result<()> { + let ctx = clickbench_test_context().await?; + let unparser = Unparser::new(&DefaultDialect {}); + + for sql in [ + r#"SELECT COUNT(DISTINCT "UserID") AS c FROM hits"#, + r#"SELECT "RegionID", COUNT(DISTINCT "UserID") AS c FROM hits GROUP BY "RegionID""#, + ] { + let optimized = ctx.sql(sql).await?.into_optimized_plan()?; + let unparsed = format!("{:#}", unparser.plan_to_sql(&optimized)?); + + let expected = sort_batches(&ctx, ctx.sql(sql).await?.collect().await?).await?; + let actual_df = ctx.sql(&unparsed).await.map_err(|e| { + datafusion_common::DataFusionError::Context( + format!("unparsed SQL failed to plan: {unparsed}"), + Box::new(e), + ) + })?; + let actual = sort_batches(&ctx, actual_df.collect().await?).await?; + + assert_eq!( + expected, actual, + "unparsed optimized plan returned different rows.\nOriginal SQL:\n{sql}\n\nUnparsed SQL:\n{unparsed}" + ); + } + + Ok(()) +} diff --git a/datafusion/sql/src/unparser/ast.rs b/datafusion/sql/src/unparser/ast.rs index eaca8c5563d91..d9788bd292e4a 100644 --- a/datafusion/sql/src/unparser/ast.rs +++ b/datafusion/sql/src/unparser/ast.rs @@ -170,6 +170,12 @@ pub struct SelectBuilder { /// Table aliases that correspond to LATERAL FLATTEN relations. /// Column references into these aliases must use `VALUE` as the column name. flatten_table_aliases: Vec, + /// Whether a `LogicalPlan::Aggregate` has already been folded into this SELECT, + /// as its select list and `GROUP BY`. A SELECT expresses at most one grouping, so + /// a second aggregate below it belongs in a derived table. + /// + /// Set with `mark_aggregated()` and read with `already_aggregated()`. + aggregated: bool, } /// Prefix used for auto-generated LATERAL FLATTEN table aliases. @@ -198,6 +204,16 @@ impl SelectBuilder { self.flatten_table_aliases.iter().any(|a| a == alias) } + /// Record that an aggregate node is now expressed by this SELECT. + pub fn mark_aggregated(&mut self) { + self.aggregated = true; + } + + /// Returns true if an aggregate node has already been folded into this SELECT. + pub fn already_aggregated(&self) -> bool { + self.aggregated + } + /// Returns the most recently generated flatten alias, or `None` if /// `next_flatten_alias` has not been called yet. pub fn current_flatten_alias(&self) -> Option { @@ -425,6 +441,7 @@ impl SelectBuilder { flavor: Some(SelectFlavor::Standard), flatten_alias_counter: 0, flatten_table_aliases: Vec::new(), + aggregated: false, } } } diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 4b385d23e941c..5da932b530183 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -884,6 +884,24 @@ impl Unparser<'_> { ) } LogicalPlan::Aggregate(agg) => { + // A SELECT expresses a single grouping, so an aggregate stacked below the + // one this SELECT already carries has to become a derived table. Stacked + // aggregates are what `single_distinct_to_groupby` produces for + // `count(DISTINCT c)`: an outer `count(alias1)` over an inner + // `GROUP BY c AS alias1`. Folding both into one SELECT would emit + // `count(alias1)` against the base table — `alias1` does not exist there, + // and where it happens to, the DISTINCT is silently gone. + if select.already_aggregated() { + return self.derive_with_dialect_alias( + "derived_aggregate", + plan, + relation, + false, + vec![], + ); + } + select.mark_aggregated(); + // Aggregation can be already handled in the projection case if !select.already_projected() { // The query returns aggregate and group expressions. If that weren't the case, diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index a9fc0716c4be2..3c1ec29ed239a 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -4274,3 +4274,72 @@ fn test_unparse_chained_intersect_build_side_is_self_contained() -> Result<()> { ); Ok(()) } + +/// Builds `count()` with the aggregate-function stub used across these tests. +fn count_col(name: &str) -> Expr { + use datafusion_expr::expr::{AggregateFunction, AggregateFunctionParams}; + Expr::AggregateFunction(AggregateFunction { + func: count_udaf(), + params: AggregateFunctionParams { + args: vec![col(name)], + distinct: false, + filter: None, + order_by: vec![], + null_treatment: None, + }, + }) +} + +/// A SELECT carries a single grouping, so an aggregate stacked directly on top of +/// another has to be unparsed as a derived table. This is the plan +/// `single_distinct_to_groupby` produces for `count(DISTINCT b)`: an outer +/// `count(alias1)` over an inner `GROUP BY b AS alias1`. Folding both into one SELECT +/// emits `count(alias1)` against the base table, where `alias1` does not exist — and +/// where a column of that name happens to exist, the DISTINCT is silently dropped. +#[test] +fn stacked_aggregate_is_unparsed_as_a_derived_table() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + ]); + + // count(DISTINCT b) — the outer aggregate groups by nothing. + let plan = table_scan(Some("test"), &schema, None)? + .aggregate(vec![col("test.b").alias("alias1")], Vec::::new())? + .aggregate(Vec::::new(), vec![count_col("alias1")])? + .project(vec![col("COUNT(alias1)").alias("count(DISTINCT test.b)")])? + .build()?; + assert_snapshot!( + plan_to_sql(&plan)?, + @r#"SELECT COUNT(alias1) AS "count(DISTINCT test.b)" FROM (SELECT test.b AS alias1 FROM test GROUP BY test.b)"# + ); + + // a, count(DISTINCT b) ... GROUP BY a — the outer aggregate keeps its own grouping, + // which must not absorb the inner one. + let plan = table_scan(Some("test"), &schema, None)? + .aggregate( + vec![col("test.a"), col("test.b").alias("alias1")], + Vec::::new(), + )? + .aggregate(vec![col("test.a")], vec![count_col("alias1")])? + .project(vec![ + col("test.a"), + col("COUNT(alias1)").alias("count(DISTINCT test.b)"), + ])? + .build()?; + assert_snapshot!( + plan_to_sql(&plan)?, + @r#"SELECT a, COUNT(alias1) AS "count(DISTINCT test.b)" FROM (SELECT test.a, test.b AS alias1 FROM test GROUP BY test.a, test.b) GROUP BY test.a"# + ); + + // A lone aggregate is still folded into the SELECT it belongs to. + let plan = table_scan(Some("test"), &schema, None)? + .aggregate(vec![col("test.a")], vec![count_col("test.b")])? + .build()?; + assert_snapshot!( + plan_to_sql(&plan)?, + @"SELECT COUNT(test.b), test.a FROM test GROUP BY test.a" + ); + + Ok(()) +}