fix: apply SingleDistinctToGroupBy to aggregates that carry an alias

datafusion/core/tests/dataframe/mod.rs

@@ -838,6 +838,84 @@ async fn aggregate_assert_no_empty_batches() -> Result<()> {
     Ok(())
 }
 
+// The DataFrame API puts an alias on the aggregate itself, where SQL leaves the
+// aggregate bare and aliases it in a projection. Both have to reach the same plan.
+#[tokio::test]
+async fn aliased_single_distinct_aggregate_is_rewritten() -> Result<()> {
+    let ctx = SessionContext::new();
+    let df = ctx
+        .read_batch(record_batch!(
+            ("a", Int32, [1, 1, 2, 2]),
+            ("b", Int32, [1, 1, 1, 2])
+        )?)?
+        .aggregate(vec![col("a")], vec![count_distinct(col("b")).alias("n")])?;
+
+    assert_snapshot!(df.clone().into_optimized_plan()?, @r"
+    Projection: ?table?.a, count(alias1) AS n
+      Aggregate: groupBy=[[?table?.a]], aggr=[[count(alias1)]]
+        Aggregate: groupBy=[[?table?.a, ?table?.b AS alias1]], aggr=[[]]
+          TableScan: ?table? projection=[a, b]
+    ");
+
+    assert_batches_eq!(
+        &[
+            "+---+---+",
+            "| a | n |",
+            "+---+---+",
+            "| 1 | 1 |",
+            "| 2 | 2 |",
+            "+---+---+",
+        ],
+        &df.sort_by(vec![col("a")])?.collect().await?
+    );
+
+    Ok(())
+}
+
+/// Stripping the aliases leaves two identical aggregates, which an `Aggregate`
+/// can't output together. They must share one, as the SQL planner's do. The
+/// `sum` follows them so its position in the outer aggregate differs from its
+/// position in the original one; `b` is `Int64` so that `sum` needs no cast,
+/// which would count as a second distinct field and skip the rewrite.
+#[tokio::test]
+async fn identical_aliased_single_distinct_aggregates_are_computed_once() -> Result<()> {
+    let ctx = SessionContext::new();
+    let df = ctx
+        .read_batch(record_batch!(
+            ("a", Int32, [1, 1, 2, 2]),
+            ("b", Int64, [1, 1, 1, 2])
+        )?)?
+        .aggregate(
+            vec![col("a")],
+            vec![
+                count_distinct(col("b")).alias("n1"),
+                count_distinct(col("b")).alias("n2"),
+                sum_distinct(col("b")).alias("s"),
+            ],
+        )?;
+
+    assert_snapshot!(df.clone().into_optimized_plan()?, @r"
+    Projection: ?table?.a, count(alias1) AS n1, count(alias1) AS n2, sum(alias1) AS s
+      Aggregate: groupBy=[[?table?.a]], aggr=[[count(alias1), sum(alias1)]]
+        Aggregate: groupBy=[[?table?.a, ?table?.b AS alias1]], aggr=[[]]
+          TableScan: ?table? projection=[a, b]
+    ");
+
+    assert_batches_eq!(
+        &[
+            "+---+----+----+---+",
+            "| a | n1 | n2 | s |",
+            "+---+----+----+---+",
+            "| 1 | 1  | 1  | 1 |",
+            "| 2 | 2  | 2  | 3 |",
+            "+---+----+----+---+",
+        ],
+        &df.sort_by(vec![col("a")])?.collect().await?
+    );
+
+    Ok(())
+}
+
 #[tokio::test]
 async fn test_aggregate_with_pk() -> Result<()> {
     // create the dataframe

datafusion/optimizer/src/eliminate_aggregate_distinct.rs

@@ -87,14 +87,15 @@ impl OptimizerRule for EliminateAggregateDistinct {
 /// `Insensitive`.
 ///
 /// If only some of them were stripped, a `Sensitive` `DISTINCT` could stay
-/// beside a stripped aggregate. A stripped aggregate carries an alias, and
-/// [`crate::single_distinct_to_groupby::SingleDistinctToGroupBy`] does not
-/// rewrite a node whose `aggr_expr` contains an alias. So `min(DISTINCT x),
-/// count(DISTINCT x)` would lose the rewrite that `count` needs. When every
-/// `DISTINCT` is `Insensitive`, none is left after stripping, so that rule has
-/// nothing to rewrite. This is conservative: `min(DISTINCT x),
-/// count(DISTINCT y)` keeps the `min` flag. That flag costs nothing at run
-/// time, because `min` ignores `is_distinct` when it selects its accumulator.
+/// beside a stripped aggregate, which
+/// [`crate::single_distinct_to_groupby::SingleDistinctToGroupBy`] then has to
+/// compute in its inner group by. While both keep the flag, `min(DISTINCT x),
+/// count(DISTINCT x)` reads the distinct column directly and needs no inner
+/// aggregate. When every `DISTINCT` is `Insensitive`, none is left after
+/// stripping, so that rule has nothing to rewrite. This is conservative:
+/// `min(DISTINCT x), count(DISTINCT y)` keeps the `min` flag. That flag costs
+/// nothing at run time, because `min` ignores `is_distinct` when it selects its
+/// accumulator.
 fn can_strip_every_distinct(aggr_expr: &[Expr]) -> Result<bool> {
     let mut found_distinct = false;
     let mut all_insensitive = true;