feat: Avoid SingleDistinctToGroupBy rule for grouped min(DISTINCT x) queries

datafusion/functions-aggregate/src/min_max.rs

@@ -46,8 +46,9 @@ use crate::min_max::min_max_bytes::MinMaxBytesAccumulator;
 use crate::min_max::min_max_struct::MinMaxStructAccumulator;
 use datafusion_common::ScalarValue;
 use datafusion_expr::{
-    Accumulator, AggregateUDFImpl, Documentation, SetMonotonicity, Signature, Volatility,
-    function::AccumulatorArgs,
+    Accumulator, AggregateUDFImpl, Documentation, Expr, SetMonotonicity, Signature,
+    Volatility,
+    function::{AccumulatorArgs, AggregateFunctionSimplification},
 };
 use datafusion_expr::{GroupsAccumulator, StatisticsArgs};
 use datafusion_macros::user_doc;
@@ -685,6 +686,18 @@ impl AggregateUDFImpl for Min {
         datafusion_expr::ReversedUDAF::Identical
     }
 
+    fn simplify(&self) -> Option<AggregateFunctionSimplification> {
+        // `min(DISTINCT x)` is identical to `min(x)`, therefore drop DISTINCT.
+        // Leave a non distinct call untouched so no rewrite is reported.
+        Some(Box::new(|mut aggregate_function, _info| {
+            if !aggregate_function.params.distinct {
+                return Ok(Expr::AggregateFunction(aggregate_function));
+            }
+            aggregate_function.params.distinct = false;
+            Ok(Expr::AggregateFunction(aggregate_function))
+        }))
+    }
+
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }

datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs

@@ -2433,6 +2433,7 @@ mod tests {
         interval_arithmetic::Interval,
         *,
     };
+    use datafusion_functions_aggregate::min_max::min_udaf;
     use datafusion_functions_window_common::field::WindowUDFFieldArgs;
     use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
     use datafusion_physical_expr::PhysicalExpr;
@@ -5245,6 +5246,37 @@ mod tests {
         assert_eq!(simplify(aggregate_function_expr), expected);
     }
 
+    /// Build `min` over `c3`, optionally distinct and/or filtered.
+    fn min_agg(distinct: bool, filter: Option<Box<Expr>>) -> Expr {
+        Expr::AggregateFunction(expr::AggregateFunction::new_udf(
+            min_udaf(),
+            vec![col("c3")],
+            distinct,
+            filter,
+            vec![],
+            None,
+        ))
+    }
+
+    #[test]
+    fn test_simplify_min_drops_distinct() {
+        // `min(DISTINCT c3)` is identically `min(c3)`
+        let (simplified, _) = simplify_with_cycle_count(min_agg(true, None));
+        assert_eq!(simplified, min_agg(false, None));
+
+        // a non distinct call is left alone and reports no change
+        let expr = min_agg(false, None);
+        let (simplified, _) = simplify_with_cycle_count(expr.clone());
+        assert_eq!(simplified, expr);
+    }
+
+    #[test]
+    fn test_simplify_min_distinct_keeps_filter() {
+        let filter = Box::new(col("c3").gt(lit(0i64)));
+        let simplified = simplify(min_agg(true, Some(filter.clone())));
+        assert_eq!(simplified, min_agg(false, Some(filter)));
+    }
+
     /// A Mock UDAF which defines `simplify` to be used in tests
     /// related to UDAF simplification
     #[derive(Debug, Clone, PartialEq, Eq, Hash)]

datafusion/optimizer/src/single_distinct_to_groupby.rs

@@ -61,6 +61,15 @@ impl SingleDistinctToGroupBy {
     }
 }
 
+fn unalias_aggregate(expr: &Expr) -> &Expr {
+    match expr {
+        Expr::Alias(alias) if matches!(*alias.expr, Expr::AggregateFunction(_)) => {
+            &alias.expr
+        }
+        _ => expr,
+    }
+}
+
 /// Check whether all aggregate exprs are distinct on a single field.
 fn is_single_distinct_agg(aggr_expr: &[Expr]) -> Result<bool> {
     let mut fields_set = HashSet::new();
@@ -76,7 +85,7 @@ fn is_single_distinct_agg(aggr_expr: &[Expr]) -> Result<bool> {
                     order_by,
                     null_treatment: _,
                 },
-        }) = expr
+        }) = unalias_aggregate(expr)
         {
             if filter.is_some() || !order_by.is_empty() {
                 return Ok(false);
@@ -179,6 +188,10 @@ impl OptimizerRule for SingleDistinctToGroupBy {
                 let mut inner_aggr_exprs = vec![];
                 let outer_aggr_exprs = aggr_expr
                     .into_iter()
+                    // The outer projection below re-aliases every aggregate back
+                    // to its original schema name, so any alias carried in here
+                    // can be dropped.
+                    .map(|aggr_expr| aggr_expr.unalias())
                     .map(|aggr_expr| match aggr_expr {
                         Expr::AggregateFunction(AggregateFunction {
                             func,

datafusion/sqllogictest/test_files/group_by.slt

@@ -4249,6 +4249,108 @@ physical_plan
 07)------------AggregateExec: mode=Partial, gby=[y@1 as y, CAST(x@0 AS Float64) as alias1], aggr=[]
 08)--------------DataSourceExec: partitions=1, partition_sizes=[1]
 
+
+statement ok
+CREATE TABLE min_distinct(g int, x int) AS VALUES
+  (1, 3), (1, 3), (1, 1), (1, NULL),
+  (2, NULL), (2, NULL),
+  (3, 7), (3, 5), (3, 5);
+
+query TT
+EXPLAIN SELECT g, min(DISTINCT x) FROM min_distinct GROUP BY g;
+----
+logical_plan
+01)Aggregate: groupBy=[[min_distinct.g]], aggr=[[min(min_distinct.x) AS min(DISTINCT min_distinct.x)]]
+02)--TableScan: min_distinct projection=[g, x]
+physical_plan
+01)AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[min(min_distinct.x) as min(DISTINCT min_distinct.x)]
+02)--RepartitionExec: partitioning=Hash([g@0], 8), input_partitions=8
+03)----AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[min(min_distinct.x) as min(DISTINCT min_distinct.x)]
+04)------RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1
+05)--------DataSourceExec: partitions=1, partition_sizes=[5]
+
+query TT
+EXPLAIN SELECT g, min(DISTINCT x) FROM min_distinct GROUP BY GROUPING SETS ((g), ());
+----
+logical_plan
+01)Projection: min_distinct.g, min(DISTINCT min_distinct.x)
+02)--Aggregate: groupBy=[[GROUPING SETS ((min_distinct.g), ())]], aggr=[[min(min_distinct.x) AS min(DISTINCT min_distinct.x)]]
+03)----TableScan: min_distinct projection=[g, x]
+physical_plan
+01)ProjectionExec: expr=[g@0 as g, min(DISTINCT min_distinct.x)@2 as min(DISTINCT min_distinct.x)]
+02)--AggregateExec: mode=FinalPartitioned, gby=[g@0 as g, __grouping_id@1 as __grouping_id], aggr=[min(min_distinct.x) as min(DISTINCT min_distinct.x)]
+03)----RepartitionExec: partitioning=Hash([g@0, __grouping_id@1], 8), input_partitions=8
+04)------AggregateExec: mode=Partial, gby=[(g@0 as g), (NULL as g)], aggr=[min(min_distinct.x) as min(DISTINCT min_distinct.x)]
+05)--------RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1
+06)----------DataSourceExec: partitions=1, partition_sizes=[5]
+
+# FILTER is preserved when the DISTINCT flag is dropped
+query TT
+EXPLAIN SELECT g, min(DISTINCT x) FILTER (WHERE x > 2) FROM min_distinct GROUP BY g;
+----
+logical_plan
+01)Aggregate: groupBy=[[min_distinct.g]], aggr=[[min(min_distinct.x) FILTER (WHERE min_distinct.x > Int32(2)) AS min(DISTINCT min_distinct.x) FILTER (WHERE min_distinct.x > Int64(2))]]
+02)--TableScan: min_distinct projection=[g, x]
+physical_plan
+01)AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[min(min_distinct.x) FILTER (WHERE min_distinct.x > Int32(2)) as min(DISTINCT min_distinct.x) FILTER (WHERE min_distinct.x > Int64(2))]
+02)--RepartitionExec: partitioning=Hash([g@0], 8), input_partitions=8
+03)----AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[min(min_distinct.x) FILTER (WHERE min_distinct.x > Int32(2)) as min(DISTINCT min_distinct.x) FILTER (WHERE min_distinct.x > Int64(2))]
+04)------RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1
+05)--------DataSourceExec: partitions=1, partition_sizes=[5]
+
+# the rule must still fire for the `count`, which is not duplicate insensitive
+query TT
+EXPLAIN SELECT g, count(DISTINCT x), min(DISTINCT x) FROM min_distinct GROUP BY g;
+----
+logical_plan
+01)Projection: min_distinct.g, count(alias1) AS count(DISTINCT min_distinct.x), min(alias2) AS min(DISTINCT min_distinct.x)
+02)--Aggregate: groupBy=[[min_distinct.g]], aggr=[[count(alias1), min(alias2)]]
+03)----Aggregate: groupBy=[[min_distinct.g, min_distinct.x AS alias1]], aggr=[[min(min_distinct.x) AS alias2]]
+04)------TableScan: min_distinct projection=[g, x]
+physical_plan
+01)ProjectionExec: expr=[g@0 as g, count(alias1)@1 as count(DISTINCT min_distinct.x), min(alias2)@2 as min(DISTINCT min_distinct.x)]
+02)--AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[count(alias1), min(alias2)]
+03)----RepartitionExec: partitioning=Hash([g@0], 8), input_partitions=8
+04)------AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[count(alias1), min(alias2)]
+05)--------AggregateExec: mode=FinalPartitioned, gby=[g@0 as g, alias1@1 as alias1], aggr=[min(min_distinct.x) as alias2]
+06)----------RepartitionExec: partitioning=Hash([g@0, alias1@1], 8), input_partitions=8
+07)------------AggregateExec: mode=Partial, gby=[g@0 as g, x@1 as alias1], aggr=[min(min_distinct.x) as alias2]
+08)--------------RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1
+09)----------------DataSourceExec: partitions=1, partition_sizes=[5]
+
+# `min(DISTINCT x)` and `min(x)` agree over duplicates, nulls and a null-only group
+query IIIII
+SELECT g, min(DISTINCT x), min(x), max(DISTINCT x), max(x)
+FROM min_distinct GROUP BY g ORDER BY g;
+----
+1 1 1 3 3
+2 NULL NULL NULL NULL
+3 5 5 7 7
+
+# a group whose every row is null
+query IIIII
+SELECT g, min(DISTINCT x), min(x), max(DISTINCT x), max(x)
+FROM min_distinct WHERE g = 2 GROUP BY g ORDER BY g;
+----
+2 NULL NULL NULL NULL
+
+# an empty input
+query II
+SELECT min(DISTINCT x), min(x) FROM min_distinct WHERE x > 100;
+----
+NULL NULL
+
+query IIII
+SELECT g, min(DISTINCT x) FILTER (WHERE x > 2), count(DISTINCT x), min(x)
+FROM min_distinct GROUP BY g ORDER BY g;
+----
+1 3 2 1
+2 NULL 0 NULL
+3 5 2 5
+
+statement ok
+DROP TABLE min_distinct;
+
 # create an unbounded table that contains ordered timestamp.
 statement ok
 CREATE UNBOUNDED EXTERNAL TABLE unbounded_csv_with_timestamps (
@@ -4432,20 +4534,20 @@ EXPLAIN SELECT c1, count(distinct c2), min(distinct c2), sum(c3), max(c4) FROM a
 ----
 logical_plan
 01)Sort: aggregate_test_100.c1 ASC NULLS LAST
-02)--Projection: aggregate_test_100.c1, count(alias1) AS count(DISTINCT aggregate_test_100.c2), min(alias1) AS min(DISTINCT aggregate_test_100.c2), sum(alias2) AS sum(aggregate_test_100.c3), max(alias3) AS max(aggregate_test_100.c4)
-03)----Aggregate: groupBy=[[aggregate_test_100.c1]], aggr=[[count(alias1), min(alias1), sum(alias2), max(alias3)]]
-04)------Aggregate: groupBy=[[aggregate_test_100.c1, aggregate_test_100.c2 AS alias1]], aggr=[[sum(CAST(aggregate_test_100.c3 AS Int64)) AS alias2, max(aggregate_test_100.c4) AS alias3]]
+02)--Projection: aggregate_test_100.c1, count(alias1) AS count(DISTINCT aggregate_test_100.c2), min(alias2) AS min(DISTINCT aggregate_test_100.c2), sum(alias3) AS sum(aggregate_test_100.c3), max(alias4) AS max(aggregate_test_100.c4)
+03)----Aggregate: groupBy=[[aggregate_test_100.c1]], aggr=[[count(alias1), min(alias2), sum(alias3), max(alias4)]]
+04)------Aggregate: groupBy=[[aggregate_test_100.c1, aggregate_test_100.c2 AS alias1]], aggr=[[min(aggregate_test_100.c2) AS alias2, sum(CAST(aggregate_test_100.c3 AS Int64)) AS alias3, max(aggregate_test_100.c4) AS alias4]]
 05)--------TableScan: aggregate_test_100 projection=[c1, c2, c3, c4]
 physical_plan
 01)SortPreservingMergeExec: [c1@0 ASC NULLS LAST]
-02)--ProjectionExec: expr=[c1@0 as c1, count(alias1)@1 as count(DISTINCT aggregate_test_100.c2), min(alias1)@2 as min(DISTINCT aggregate_test_100.c2), sum(alias2)@3 as sum(aggregate_test_100.c3), max(alias3)@4 as max(aggregate_test_100.c4)]
+02)--ProjectionExec: expr=[c1@0 as c1, count(alias1)@1 as count(DISTINCT aggregate_test_100.c2), min(alias2)@2 as min(DISTINCT aggregate_test_100.c2), sum(alias3)@3 as sum(aggregate_test_100.c3), max(alias4)@4 as max(aggregate_test_100.c4)]
 03)----SortExec: expr=[c1@0 ASC NULLS LAST], preserve_partitioning=[true]
-04)------AggregateExec: mode=FinalPartitioned, gby=[c1@0 as c1], aggr=[count(alias1), min(alias1), sum(alias2), max(alias3)]
+04)------AggregateExec: mode=FinalPartitioned, gby=[c1@0 as c1], aggr=[count(alias1), min(alias2), sum(alias3), max(alias4)]
 05)--------RepartitionExec: partitioning=Hash([c1@0], 8), input_partitions=8
-06)----------AggregateExec: mode=Partial, gby=[c1@0 as c1], aggr=[count(alias1), min(alias1), sum(alias2), max(alias3)]
-07)------------AggregateExec: mode=FinalPartitioned, gby=[c1@0 as c1, alias1@1 as alias1], aggr=[sum(aggregate_test_100.c3) as alias2, max(aggregate_test_100.c4) as alias3]
+06)----------AggregateExec: mode=Partial, gby=[c1@0 as c1], aggr=[count(alias1), min(alias2), sum(alias3), max(alias4)]
+07)------------AggregateExec: mode=FinalPartitioned, gby=[c1@0 as c1, alias1@1 as alias1], aggr=[min(aggregate_test_100.c2) as alias2, sum(aggregate_test_100.c3) as alias3, max(aggregate_test_100.c4) as alias4]
 08)--------------RepartitionExec: partitioning=Hash([c1@0, alias1@1], 8), input_partitions=8
-09)----------------AggregateExec: mode=Partial, gby=[c1@0 as c1, c2@1 as alias1], aggr=[sum(aggregate_test_100.c3) as alias2, max(aggregate_test_100.c4) as alias3]
+09)----------------AggregateExec: mode=Partial, gby=[c1@0 as c1, c2@1 as alias1], aggr=[min(aggregate_test_100.c2) as alias2, sum(aggregate_test_100.c3) as alias3, max(aggregate_test_100.c4) as alias4]
 10)------------------RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1
 11)--------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100.csv]]}, projection=[c1, c2, c3, c4], file_type=csv, has_header=true