feat: DISTINCT handling for Aggregate UDFs

datafusion/core/src/optimizer_rule_reference.md

@@ -56,12 +56,13 @@ Rule order matters. The default pipeline may change between releases.
 | 17    | `eliminate_outer_join`                    | Rewrites outer joins to inner joins when later filters reject the NULL-extended rows.                                       |
 | 18    | `push_down_limit`                         | Moves literal limits closer to scans and unions and merges adjacent limits.                                                 |
 | 19    | `push_down_filter`                        | Moves filters as early as possible through filter-commutative operators.                                                    |
-| 20    | `single_distinct_aggregation_to_group_by` | Rewrites single-column `DISTINCT` aggregations into two-stage `GROUP BY` plans.                                             |
-| 21    | `eliminate_group_by_constant`             | Removes constant or functionally redundant expressions from `GROUP BY`.                                                     |
-| 22    | `common_sub_expression_eliminate`         | Computes repeated subexpressions once and reuses the result.                                                                |
-| 23    | `extract_leaf_expressions`                | Pulls cheap leaf expressions closer to data sources so later pruning and filter rules can act earlier.                      |
-| 24    | `push_down_leaf_projections`              | Pushes the helper projections created by leaf extraction toward leaf inputs.                                                |
-| 25    | `optimize_projections`                    | Prunes unused columns and removes unnecessary logical projections.                                                          |
+| 20    | `eliminate_aggregate_distinct`            | Drops the `DISTINCT` modifier from aggregates whose result cannot change, such as `min`, `max` and `bit_or`.                |
+| 21    | `single_distinct_aggregation_to_group_by` | Rewrites single-column `DISTINCT` aggregations into two-stage `GROUP BY` plans.                                             |
+| 22    | `eliminate_group_by_constant`             | Removes constant or functionally redundant expressions from `GROUP BY`.                                                     |
+| 23    | `common_sub_expression_eliminate`         | Computes repeated subexpressions once and reuses the result.                                                                |
+| 24    | `extract_leaf_expressions`                | Pulls cheap leaf expressions closer to data sources so later pruning and filter rules can act earlier.                      |
+| 25    | `push_down_leaf_projections`              | Pushes the helper projections created by leaf extraction toward leaf inputs.                                                |
+| 26    | `optimize_projections`                    | Prunes unused columns and removes unnecessary logical projections.                                                          |
 
 ### Physical Optimizer Rules
 

datafusion/expr/src/lib.rs

@@ -137,10 +137,10 @@ pub use partition_evaluator::PartitionEvaluator;
 pub use sqlparser;
 pub use table_source::{TableProviderFilterPushDown, TableSource, TableType};
 pub use udaf::{
-    AggregateUDF, AggregateUDFImpl, ReversedUDAF, SetMonotonicity, StatisticsArgs,
-    UdafDisplayNameBuilder, UdafHumanDisplayBuilder, UdafSchemaNameBuilder,
-    UdafWindowFunctionDisplayNameBuilder, UdafWindowFunctionSchemaNameBuilder,
-    udaf_default_return_field,
+    AggregateUDF, AggregateUDFImpl, DistinctHandling, ReversedUDAF, SetMonotonicity,
+    StatisticsArgs, UdafDisplayNameBuilder, UdafHumanDisplayBuilder,
+    UdafSchemaNameBuilder, UdafWindowFunctionDisplayNameBuilder,
+    UdafWindowFunctionSchemaNameBuilder, udaf_default_return_field,
 };
 #[expect(deprecated)]
 pub use udaf::{

datafusion/expr/src/udaf.rs

@@ -360,6 +360,11 @@ impl AggregateUDF {
         self.inner.supports_within_group_clause()
     }
 
+    /// See [`AggregateUDFImpl::distinct_handling`] for more details.
+    pub fn distinct_handling(&self) -> DistinctHandling {
+        self.inner.distinct_handling()
+    }
+
     /// Returns the documentation for this Aggregate UDF.
     ///
     /// Documentation can be accessed programmatically as well as
@@ -940,6 +945,20 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any {
         false
     }
 
+    /// How this function treats the `DISTINCT` modifier.
+    ///
+    /// Return [`DistinctHandling::Ignored`] for duplicate-insensitive
+    /// functions so that `f(DISTINCT x)` is planned as `f(x)`.
+    ///
+    /// Return [`DistinctHandling::Unsupported`] if the accumulator does not
+    /// implement `DISTINCT`, that is, it does not read `is_distinct`, or it
+    /// rejects `DISTINCT` with an error. The planner then has to deduplicate
+    /// the input or reject the query. Nothing reads this variant yet:
+    /// rejecting such queries at planning time is a follow-up change.
+    fn distinct_handling(&self) -> DistinctHandling {
+        DistinctHandling::Honored
+    }
+
     /// Returns the documentation for this Aggregate UDF.
     ///
     /// Documentation can be accessed programmatically as well as
@@ -1687,6 +1706,10 @@ impl AggregateUDFImpl for AliasedAggregateUDFImpl {
         self.inner.set_monotonicity(data_type)
     }
 
+    fn distinct_handling(&self) -> DistinctHandling {
+        self.inner.distinct_handling()
+    }
+
     fn documentation(&self) -> Option<&Documentation> {
         self.inner.documentation()
     }
@@ -1713,6 +1736,29 @@ pub enum SetMonotonicity {
     NotMonotonic,
 }
 
+/// How an aggregate function treats the `DISTINCT` modifier.
+///
+/// Mathematically, `Ignored` means the function's merge operation is
+/// idempotent (its state forms a semilattice): f(S ⊎ S) = f(S), so
+/// removing duplicates from the input cannot change the result.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[non_exhaustive]
+pub enum DistinctHandling {
+    /// The result is the same with or without `DISTINCT`, so the planner
+    /// is free to drop it. `min`, `max`, `bool_and`, `bit_or`, ...
+    Ignored,
+    /// The accumulator reads `AccumulatorArgs::is_distinct` and deduplicates
+    /// its input, so the planner must leave the flag alone. `count`, `sum`,
+    /// `avg`, `var_samp`, `array_agg`, ... This is the default.
+    Honored,
+    /// The accumulator does not implement `DISTINCT`: it does not read
+    /// `is_distinct`, or it rejects `DISTINCT` with an error. The planner has
+    /// to deduplicate the input first (today `SingleDistinctToGroupBy` does
+    /// that for single-argument functions) or reject the query. `stddev`,
+    /// `approx_median`, `corr`, `regr_*`, `nth_value`, ...
+    Unsupported,
+}
+
 #[cfg(test)]
 mod test {
     use crate::{AggregateUDF, AggregateUDFImpl};

datafusion/functions-aggregate/src/any_value.rs

@@ -122,4 +122,9 @@ impl AggregateUDFImpl for AnyValue {
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
+
+    // TODO: this is arguably `DistinctHandling::Ignored` — the accumulator
+    // ignores `is_distinct` and returns an unspecified input value either
+    // way. Grouped with `first_value`/`last_value` and left at the default
+    // `Honored` until that family is settled together.
 }

datafusion/functions-aggregate/src/approx_distinct.rs

@@ -38,6 +38,7 @@ use datafusion_common::{
     DataFusionError, Result, downcast_value, internal_datafusion_err, internal_err,
     not_impl_err,
 };
+use datafusion_expr::DistinctHandling;
 use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
 use datafusion_expr::utils::format_state_name;
 use datafusion_expr::{
@@ -871,6 +872,11 @@ impl AggregateUDFImpl for ApproxDistinct {
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
+
+    fn distinct_handling(&self) -> DistinctHandling {
+        // Updating an HLL register with a value already seen is a no-op.
+        DistinctHandling::Ignored
+    }
 }
 
 fn is_fixed_domain_type(data_type: &DataType) -> bool {

datafusion/functions-aggregate/src/approx_median.rs

@@ -147,4 +147,10 @@ impl AggregateUDFImpl for ApproxMedian {
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
+
+    fn distinct_handling(&self) -> datafusion_expr::DistinctHandling {
+        // The accumulator rejects `DISTINCT` with `not_impl_err!`, so the
+        // planner has to deduplicate the input first.
+        datafusion_expr::DistinctHandling::Unsupported
+    }
 }

datafusion/functions-aggregate/src/approx_percentile_cont.rs

@@ -31,6 +31,7 @@ use datafusion_common::{
     DataFusionError, Result, ScalarValue, downcast_value, internal_err, not_impl_err,
     plan_err,
 };
+use datafusion_expr::DistinctHandling;
 use datafusion_expr::expr::{AggregateFunction, Sort};
 use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
 use datafusion_expr::utils::format_state_name;
@@ -324,6 +325,13 @@ impl AggregateUDFImpl for ApproxPercentileCont {
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
+
+    fn distinct_handling(&self) -> DistinctHandling {
+        // Duplicate-sensitive, but the accumulator does not read
+        // `is_distinct` and today silently returns the non-distinct answer.
+        // The tag records the intent; enforcement is a follow-up change.
+        DistinctHandling::Unsupported
+    }
 }
 
 #[derive(Debug)]

datafusion/functions-aggregate/src/approx_percentile_cont_with_weight.rs

@@ -26,6 +26,7 @@ use arrow::{array::ArrayRef, datatypes::DataType};
 use datafusion_common::ScalarValue;
 use datafusion_common::types::{NativeType, logical_float64};
 use datafusion_common::{Result, not_impl_err, plan_err};
+use datafusion_expr::DistinctHandling;
 use datafusion_expr::expr::{AggregateFunction, Sort};
 use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
 use datafusion_expr::{
@@ -281,6 +282,11 @@ impl AggregateUDFImpl for ApproxPercentileContWithWeight {
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
+
+    fn distinct_handling(&self) -> DistinctHandling {
+        // The accumulator rejects `DISTINCT` with `not_impl_err!`.
+        DistinctHandling::Unsupported
+    }
 }
 
 #[derive(Debug)]

datafusion/functions-aggregate/src/bit_and_or_xor.rs

@@ -31,6 +31,7 @@ use datafusion_common::hash_utils::RandomState;
 
 use datafusion_common::cast::as_list_array;
 use datafusion_common::{Result, ScalarValue, not_impl_err};
+use datafusion_expr::DistinctHandling;
 use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
 use datafusion_expr::utils::format_state_name;
 use datafusion_expr::{
@@ -318,6 +319,18 @@ impl AggregateUDFImpl for BitwiseOperation {
     fn documentation(&self) -> Option<&Documentation> {
         Some(self.documentation)
     }
+
+    fn distinct_handling(&self) -> DistinctHandling {
+        match self.operation {
+            // Bitwise AND/OR are idempotent: duplicates cannot change the
+            // result. Only XOR has a distinct accumulator.
+            BitwiseOperationType::And | BitwiseOperationType::Or => {
+                DistinctHandling::Ignored
+            }
+            // XOR cancels duplicate pairs, so `DISTINCT` is meaningful.
+            BitwiseOperationType::Xor => DistinctHandling::Honored,
+        }
+    }
 }
 
 struct BitAndAccumulator<T: ArrowNumericType> {

datafusion/functions-aggregate/src/bool_and_or.rs

@@ -29,6 +29,7 @@ use arrow::datatypes::{DataType, FieldRef};
 use datafusion_common::internal_err;
 use datafusion_common::{Result, ScalarValue};
 use datafusion_common::{downcast_value, not_impl_err};
+use datafusion_expr::DistinctHandling;
 use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
 use datafusion_expr::utils::{AggregateOrderSensitivity, format_state_name};
 use datafusion_expr::{
@@ -183,6 +184,11 @@ impl AggregateUDFImpl for BoolAnd {
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
+
+    fn distinct_handling(&self) -> DistinctHandling {
+        // Boolean AND/OR are idempotent: duplicates cannot change the result.
+        DistinctHandling::Ignored
+    }
 }
 
 #[derive(Debug, Default)]
@@ -313,6 +319,11 @@ impl AggregateUDFImpl for BoolOr {
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
+
+    fn distinct_handling(&self) -> DistinctHandling {
+        // Boolean AND/OR are idempotent: duplicates cannot change the result.
+        DistinctHandling::Ignored
+    }
 }
 
 #[derive(Debug, Default)]

datafusion/functions-aggregate/src/correlation.rs

@@ -31,6 +31,7 @@ use arrow::{
     array::ArrayRef,
     datatypes::{DataType, Field},
 };
+use datafusion_expr::DistinctHandling;
 use datafusion_expr::{EmitTo, GroupSelection, GroupsAccumulator};
 use datafusion_functions_aggregate_common::aggregate::groups_accumulator::accumulate::accumulate_multiple;
 use log::debug;
@@ -145,6 +146,13 @@ impl AggregateUDFImpl for Correlation {
         debug!("GroupsAccumulator is created for aggregate function `corr(c1, c2)`");
         Ok(Box::new(CorrelationGroupsAccumulator::new()))
     }
+
+    fn distinct_handling(&self) -> DistinctHandling {
+        // Duplicate-sensitive, but the accumulator does not read
+        // `is_distinct` and today silently returns the non-distinct answer.
+        // The tag records the intent; enforcement is a follow-up change.
+        DistinctHandling::Unsupported
+    }
 }
 
 /// An accumulator to compute correlation

datafusion/functions-aggregate/src/covariance.rs

@@ -21,6 +21,7 @@ use arrow::array::ArrayRef;
 use arrow::datatypes::{DataType, Field, FieldRef};
 use datafusion_common::cast::{as_float64_array, as_uint64_array};
 use datafusion_common::{Result, ScalarValue};
+use datafusion_expr::DistinctHandling;
 use datafusion_expr::{
     Accumulator, AggregateUDFImpl, Documentation, Signature, Volatility,
     function::{AccumulatorArgs, StateFieldsArgs},
@@ -128,6 +129,13 @@ impl AggregateUDFImpl for CovarianceSample {
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
+
+    fn distinct_handling(&self) -> DistinctHandling {
+        // Duplicate-sensitive, but the accumulator does not read
+        // `is_distinct` and today silently returns the non-distinct answer.
+        // The tag records the intent; enforcement is a follow-up change.
+        DistinctHandling::Unsupported
+    }
 }
 
 #[user_doc(
@@ -206,6 +214,13 @@ impl AggregateUDFImpl for CovariancePopulation {
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
+
+    fn distinct_handling(&self) -> DistinctHandling {
+        // Duplicate-sensitive, but the accumulator does not read
+        // `is_distinct` and today silently returns the non-distinct answer.
+        // The tag records the intent; enforcement is a follow-up change.
+        DistinctHandling::Unsupported
+    }
 }
 
 /// An accumulator to compute covariance

datafusion/functions-aggregate/src/first_last.rs

@@ -380,6 +380,12 @@ impl AggregateUDFImpl for FirstValue {
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
+
+    // TODO: whether this is `DistinctHandling::Ignored` depends on `ORDER BY`.
+    // `first_value(DISTINCT x ORDER BY y)` deduplicates `x` and leaves the `y`
+    // ordering meaningless, while `first_value(DISTINCT x ORDER BY x)` is just
+    // `min(x)`. Left at the default `Honored` until that is settled, even
+    // though the accumulator ignores `is_distinct` today.
 }
 
 struct FirstLastGroupsAccumulator<S: ValueState> {
@@ -1294,6 +1300,12 @@ impl AggregateUDFImpl for LastValue {
     ) -> Result<Box<dyn GroupsAccumulator>> {
         create_groups_accumulator(&args, false, self.is_input_pre_ordered, self.name())
     }
+
+    // TODO: whether this is `DistinctHandling::Ignored` depends on `ORDER BY`.
+    // `last_value(DISTINCT x ORDER BY y)` deduplicates `x` and leaves the `y`
+    // ordering meaningless, while `last_value(DISTINCT x ORDER BY x)` is
+    // `max(x)` when `x` has no NULL. Left at the default `Honored` until that is settled, even
+    // though the accumulator ignores `is_distinct` today.
 }
 
 /// This accumulator is used when there is no ordering specified for the

datafusion/functions-aggregate/src/grouping.rs

@@ -20,6 +20,7 @@
 use arrow::datatypes::Field;
 use arrow::datatypes::{DataType, FieldRef};
 use datafusion_common::{Result, not_impl_err};
+use datafusion_expr::DistinctHandling;
 use datafusion_expr::function::AccumulatorArgs;
 use datafusion_expr::function::StateFieldsArgs;
 use datafusion_expr::utils::format_state_name;
@@ -110,4 +111,13 @@ impl AggregateUDFImpl for Grouping {
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
+
+    fn distinct_handling(&self) -> DistinctHandling {
+        // The result depends only on which grouping set a row belongs to, not
+        // on how many rows share a value, so duplicates cannot change it.
+        // `ResolveGroupingFunction` replaces the call before the optimizer
+        // runs, so this tag is not reachable from SQL and the accumulator
+        // above is never built.
+        DistinctHandling::Ignored
+    }
 }

datafusion/functions-aggregate/src/min_max.rs

@@ -45,6 +45,7 @@ use arrow::datatypes::{
 use crate::min_max::min_max_bytes::MinMaxBytesAccumulator;
 use crate::min_max::min_max_struct::MinMaxStructAccumulator;
 use datafusion_common::ScalarValue;
+use datafusion_expr::DistinctHandling;
 use datafusion_expr::{
     Accumulator, AggregateUDFImpl, Documentation, SetMonotonicity, Signature, Volatility,
     function::AccumulatorArgs,
@@ -399,6 +400,11 @@ impl AggregateUDFImpl for Max {
         // the same as new values are seen.
         SetMonotonicity::Increasing
     }
+
+    fn distinct_handling(&self) -> DistinctHandling {
+        // `MAX` is idempotent: duplicates cannot change the maximum.
+        DistinctHandling::Ignored
+    }
 }
 
 #[derive(Debug)]
@@ -694,6 +700,11 @@ impl AggregateUDFImpl for Min {
         // the same as new values are seen.
         SetMonotonicity::Decreasing
     }
+
+    fn distinct_handling(&self) -> DistinctHandling {
+        // `MIN` is idempotent: duplicates cannot change the minimum.
+        DistinctHandling::Ignored
+    }
 }
 
 #[derive(Debug)]