fix(unparser): unparse a stacked aggregate as a derived table

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.

Author: claudespice

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(())
+}

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<String>,
+    /// 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<String> {
@@ -425,6 +441,7 @@ impl SelectBuilder {
             flavor: Some(SelectFlavor::Standard),
             flatten_alias_counter: 0,
             flatten_table_aliases: Vec::new(),
+            aggregated: false,
         }
     }
 }

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,

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(<col>)` 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::<Expr>::new())?
+        .aggregate(Vec::<Expr>::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::<Expr>::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(())
+}