fix: keep a DISTINCT aggregate's arity in SingleDistinctToGroupBy (#25461)

## Which issue does this PR close?

- Closes #25417.

## Rationale for this change

`is_single_distinct_agg` requires that every distinct argument in the
aggregation is the same expression. It checks that by collecting the
arguments into a set and requiring the set to hold exactly one element:

```rust
if *distinct {
    for e in args {
        fields_set.insert(e);
    }
```

A call that passes that expression more than once satisfies the check,
because its arguments collapse to a single entry. `corr(DISTINCT x, x)`
is such a call, and the rewrite below then assumed one argument per
distinct aggregate and returned an internal error.

This is also why the neighbouring cases in the issue are unaffected:
`corr(DISTINCT x, y)` and `corr(DISTINCT x, x + 0.0)` put two entries in
the set, so the rule bails out early and leaves the plan alone.

## What changes are included in this PR?

The distinct branch keeps the call's arity, repeating the alias that the
inner group by produces, so `corr(DISTINCT x, x)` becomes a group by `x`
with `corr(alias1, alias1)` above it.

That is sound because of the check above. Every distinct argument is the
same expression, so the distinct argument tuples the call aggregates are
exactly the distinct values of that expression, which is what grouping
by it produces.

The assertion stays, now stating the invariant the rewrite actually
relies on rather than an arity the rule does not enforce. It also covers
an empty argument list, which would otherwise panic in
`args.swap_remove(0)` instead of reporting an internal error.

## What is the testing strategy for this PR?

Rule tests in `single_distinct_to_groupby.rs`:

- `single_distinct_repeated_arg_and_groupby`: `corr(DISTINCT b, b)` is
rewritten, and the plan keeps both arguments and the original output
column name.
- `single_distinct_two_args_and_groupby`: `corr(DISTINCT b, c)` is still
left alone.

Two queries in `aggregate.slt` cover it end to end. One is `covar_samp`,
chosen because its value depends on the de-duplication rather than only
on the absence of the error: over the same input it returns `0.5` with
`DISTINCT` and `0.333333333333` without, so the expected value pins the
semantics.

Measured, rather than assumed:

| check | result |
| --- | --- |
| both halves of the fix reverted | the two `aggregate.slt` queries and
`single_distinct_repeated_arg_and_groupby` fail with an `Assertion
failed: args.len() == 1` internal error, which is the failure in the
issue |
| the rewrite against a reference the rule does not touch | the same
aggregates over a manually de-duplicated input return the same values |

`cargo test -p datafusion-optimizer` (915 tests), the sqllogictest suite
(521 files), and `./ci/scripts/rust_clippy.sh` pass.

## Are there any user-facing changes?

Queries that failed with an internal error now run. No plan that the
rule rewrote before is rewritten differently: a single argument call
produces the same outer aggregate as it did, and the only calls that
reach this branch with more than one argument are the ones that
previously returned the error.

Author: namanjain24-sudo

datafusion/optimizer/src/single_distinct_to_groupby.rs

@@ -23,7 +23,7 @@ use crate::optimizer::ApplyOrder;
 use crate::{OptimizerConfig, OptimizerRule};
 
 use datafusion_common::{
-    DFSchema, DataFusionError, HashSet, Result, assert_eq_or_internal_err,
+    DFSchema, DataFusionError, HashSet, Result, assert_or_internal_err,
     tree_node::Transformed,
 };
 use datafusion_expr::builder::project;
@@ -314,11 +314,18 @@ impl OptimizerRule for SingleDistinctToGroupBy {
                                 },
                         }) => {
                             if distinct {
-                                assert_eq_or_internal_err!(
-                                    args.len(),
-                                    1,
-                                    "DISTINCT aggregate should have exactly one argument"
+                                // `is_single_distinct_agg` checked that every
+                                // distinct argument in the plan is the same
+                                // expression. A call may still take it more than
+                                // once, as `corr(DISTINCT x, x)` does, and keeps
+                                // its arity here: its distinct argument tuples
+                                // are the distinct values of that expression.
+                                assert_or_internal_err!(
+                                    args.first()
+                                        .is_some_and(|f| args.iter().all(|a| a == f)),
+                                    "DISTINCT aggregate arguments should be the same expression"
                                 );
+                                let arity = args.len();
                                 let arg = args.swap_remove(0);
 
                                 if group_fields_set.insert(arg.schema_name().to_string())
@@ -329,7 +336,7 @@ impl OptimizerRule for SingleDistinctToGroupBy {
                                 let outer =
                                     Expr::AggregateFunction(AggregateFunction::new_udf(
                                         func,
-                                        vec![col(SINGLE_DISTINCT_ALIAS)],
+                                        vec![col(SINGLE_DISTINCT_ALIAS); arity],
                                         false, // intentional to remove distinct here
                                         filter,
                                         order_by,
@@ -449,6 +456,7 @@ mod tests {
         lit,
         logical_plan::builder::{LogicalPlanBuilder, table_scan},
     };
+    use datafusion_functions_aggregate::correlation::corr_udaf;
     use datafusion_functions_aggregate::count::count_udaf;
     use datafusion_functions_aggregate::expr_fn::{count, count_distinct, max, min, sum};
     use datafusion_functions_aggregate::min_max::max_udaf;
@@ -465,6 +473,17 @@ mod tests {
         ))
     }
 
+    fn corr_distinct(y: Expr, x: Expr) -> Expr {
+        Expr::AggregateFunction(AggregateFunction::new_udf(
+            corr_udaf(),
+            vec![y, x],
+            true,
+            None,
+            vec![],
+            None,
+        ))
+    }
+
     fn sum_distinct(expr: Expr) -> Expr {
         Expr::AggregateFunction(AggregateFunction::new_udf(
             sum_udaf(),
@@ -770,6 +789,46 @@ mod tests {
         )
     }
 
+    #[test]
+    fn single_distinct_repeated_arg_and_groupby() -> Result<()> {
+        let table_scan = test_table_scan()?;
+
+        let plan = LogicalPlanBuilder::from(table_scan)
+            .aggregate(vec![col("a")], vec![corr_distinct(col("b"), col("b"))])?
+            .build()?;
+
+        // Should work: the call takes the distinct argument twice, so it keeps
+        // both arguments after the rewrite.
+        assert_optimized_plan_equal!(
+            plan,
+            @r"
+        Projection: test.a, corr(alias1,alias1) AS corr(DISTINCT test.b,test.b) [a:UInt32, corr(DISTINCT test.b,test.b):Float64;N]
+          Aggregate: groupBy=[[test.a]], aggr=[[corr(alias1, alias1)]] [a:UInt32, corr(alias1,alias1):Float64;N]
+            Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[]] [a:UInt32, alias1:UInt32]
+              TableScan: test [a:UInt32, b:UInt32, c:UInt32]
+        "
+        )
+    }
+
+    #[test]
+    fn single_distinct_two_args_and_groupby() -> Result<()> {
+        let table_scan = test_table_scan()?;
+
+        let plan = LogicalPlanBuilder::from(table_scan)
+            .aggregate(vec![col("a")], vec![corr_distinct(col("b"), col("c"))])?
+            .build()?;
+
+        // Do nothing: the arguments are different expressions, so there is no
+        // single field to group by.
+        assert_optimized_plan_equal!(
+            plan,
+            @r"
+        Aggregate: groupBy=[[test.a]], aggr=[[corr(DISTINCT test.b, test.c)]] [a:UInt32, corr(DISTINCT test.b,test.c):Float64;N]
+          TableScan: test [a:UInt32, b:UInt32, c:UInt32]
+        "
+        )
+    }
+
     #[test]
     fn one_field_two_distinct_and_groupby() -> Result<()> {
         let table_scan = test_table_scan()?;

datafusion/sqllogictest/test_files/aggregate.slt

@@ -404,6 +404,34 @@ GROUP BY g
 1 1
 2 NULL
 
+# A DISTINCT aggregate may take the distinct argument more than once.
+# Group 1 aggregates over the distinct values 1.0 and 2.0, group 2 over the
+# single distinct value 3.0, which leaves the sample covariance undefined.
+# Group 3 repeats a NULL row alongside the distinct value 2.0: DISTINCT
+# collapses the repeated NULLs into one, and the aggregate then excludes
+# that NULL the same way it would without the rewrite, leaving a single
+# distinct point - again an undefined sample covariance. Group 4 is all
+# NULL, so the aggregate has zero non-null distinct points to work with.
+query IR rowsort
+SELECT g, corr(DISTINCT x, x)
+FROM (VALUES (1, 1.0), (1, 1.0), (1, 2.0), (2, 3.0), (2, 3.0), (3, NULL), (3, NULL), (3, 2.0), (4, NULL), (4, NULL)) AS t(g, x)
+GROUP BY g
+----
+1 1
+2 NULL
+3 NULL
+4 NULL
+
+query IR rowsort
+SELECT g, covar_samp(DISTINCT x, x)
+FROM (VALUES (1, 1.0), (1, 1.0), (1, 2.0), (2, 3.0), (2, 3.0), (3, NULL), (3, NULL), (3, 2.0), (4, NULL), (4, NULL)) AS t(g, x)
+GROUP BY g
+----
+1 0.5
+2 NULL
+3 NULL
+4 NULL
+
 # all_nulls_query_correlation
 query R
 with data as (