perf: Optimize min(DISTINCT x) aggregations

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,14 @@ impl AggregateUDFImpl for Min {
         datafusion_expr::ReversedUDAF::Identical
     }
 
+    fn simplify(&self) -> Option<AggregateFunctionSimplification> {
+        // `min(DISTINCT x)` is identical to `min(x)`, therefore drop DISTINCT.
+        Some(Box::new(|mut aggregate_function, _info| {
+            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

@@ -2489,6 +2489,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;
@@ -5395,6 +5396,23 @@ mod tests {
         assert_eq!(simplify(aggregate_function_expr), expected);
     }
 
+    #[test]
+    fn test_simplify_min_drops_distinct() {
+        let min_agg = |distinct: bool| {
+            Expr::AggregateFunction(expr::AggregateFunction::new_udf(
+                min_udaf(),
+                vec![col("c3")],
+                distinct,
+                None,
+                vec![],
+                None,
+            ))
+        };
+
+        let simplified = simplify(min_agg(true));
+        assert_eq!(simplified, min_agg(false));
+    }
+
     /// 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

@@ -120,6 +120,15 @@ impl CountRollup {
     }
 }
 
+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],
@@ -141,7 +150,7 @@ fn is_single_distinct_agg(
                     order_by,
                     null_treatment: _,
                 },
-        }) = expr
+        }) = unalias_aggregate(expr)
         {
             if filter.is_some() || !order_by.is_empty() {
                 return Ok(false);
@@ -301,6 +310,7 @@ impl OptimizerRule for SingleDistinctToGroupBy {
                 // zero that `sum` reports as NULL over an empty input.
                 let (outer_aggr_exprs, outer_proj_exprs): (Vec<Expr>, Vec<Expr>) = aggr_expr
                     .into_iter()
+                    .map(|aggr_expr| aggr_expr.unalias())
                     .map(|aggr_expr| match aggr_expr {
                         Expr::AggregateFunction(AggregateFunction {
                             func,
@@ -321,7 +331,7 @@ impl OptimizerRule for SingleDistinctToGroupBy {
                                 );
                                 let arg = args.swap_remove(0);
 
-                                if group_fields_set.insert(arg.schema_name().to_string())
+                                  if group_fields_set.insert(arg.schema_name().to_string())
                                 {
                                     inner_group_exprs
                                         .push(arg.alias(SINGLE_DISTINCT_ALIAS));

datafusion/sqllogictest/test_files/group_by.slt

@@ -4249,6 +4249,29 @@ 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]
+
+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 +4455,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
 

datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt

@@ -131,7 +131,7 @@ EXPLAIN SELECT g, count(*) AS records, min(DISTINCT v) AS distinct_min_v FROM t
 ----
 logical_plan
 01)Projection: t.g, count(Int64(1)) AS count(*) AS records, min(DISTINCT t.v) AS distinct_min_v
-02)--Aggregate: groupBy=[[t.g]], aggr=[[count(Int64(1)), min(DISTINCT t.v)]]
+02)--Aggregate: groupBy=[[t.g]], aggr=[[count(Int64(1)), min(t.v) AS min(DISTINCT t.v)]]
 03)----TableScan: t projection=[g, v]
 
 # The gate covers only the count. A plan that already qualified through sum,