feat: DISTINCT handling for Aggregate UDFs

datafusion/core/src/optimizer_rule_reference.md

@@ -35,33 +35,34 @@ Rule order matters. The default pipeline may change between releases.
 
 ### Logical Optimizer Rules
 
-| order | rule                                      | summary                                                                                                                                                               |
-| ----- | ----------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| 1     | `rewrite_set_comparison`                  | Rewrites `ANY` and `ALL` set-comparison subqueries into `EXISTS`-based boolean expressions with correct SQL NULL semantics.                                           |
-| 2     | `optimize_unions`                         | Flattens nested unions and removes unions with a single input.                                                                                                        |
-| 3     | `unions_to_filter`                        | Merges `UNION DISTINCT` branches that share the same source into a single filtered branch with a disjunctive predicate.                                               |
-| 4     | `simplify_expressions`                    | Constant-folds and simplifies expressions while preserving output names.                                                                                              |
-| 5     | `replace_distinct_aggregate`              | Rewrites `DISTINCT` and `DISTINCT ON` operators into aggregate-based plans that later rules can optimize further.                                                     |
-| 6     | `eliminate_join`                          | Replaces keyless inner joins with a literal `false` filter by an empty relation.                                                                                      |
-| 7     | `decorrelate_predicate_subquery`          | Converts eligible `IN` and `EXISTS` predicate subqueries into semi or anti joins.                                                                                     |
-| 8     | `scalar_subquery_to_join`                 | Rewrites eligible scalar subqueries into joins and adds schema-preserving projections.                                                                                |
-| 9     | `decorrelate_lateral_join`                | Rewrites eligible lateral joins into regular joins.                                                                                                                   |
-| 10    | `extract_equijoin_predicate`              | Splits join filters into equijoin keys and residual predicates.                                                                                                       |
-| 11    | `eliminate_duplicated_expr`               | Removes duplicate expressions from projections, aggregates, and similar operators.                                                                                    |
-| 12    | `eliminate_filter`                        | Drops always-true filters and replaces always-false or NULL filters with empty relations.                                                                             |
-| 13    | `eliminate_cross_join`                    | Uses filter predicates to replace cross joins with inner joins when join keys can be found.                                                                           |
-| 14    | `eliminate_limit`                         | Removes no-op limits and simplifies trivial limit shapes.                                                                                                             |
-| 15    | `propagate_empty_relation`                | Pushes empty-relation knowledge upward so operators fed by no rows collapse early.                                                                                    |
-| 16    | `filter_null_join_keys`                   | Adds `IS NOT NULL` filters to nullable equijoin keys that can never match.                                                                                            |
-| 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 to scans/unions, merges limits, and pushes `Sort(fetch=N)` onto a preserved join side only when every sort key is from it; the outer Sort stays. |
-| 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.                                                                                                    |
+| order | rule                                      | summary                                                                                                                     |
+|-------| ----------------------------------------- | --------------------------------------------------------------------------------------------------------------------------- |
+| 1     | `rewrite_set_comparison`                  | Rewrites `ANY` and `ALL` set-comparison subqueries into `EXISTS`-based boolean expressions with correct SQL NULL semantics. |
+| 2     | `optimize_unions`                         | Flattens nested unions and removes unions with a single input.                                                              |
+| 3     | `unions_to_filter`                        | Merges `UNION DISTINCT` branches that share the same source into a single filtered branch with a disjunctive predicate.     |
+| 4     | `simplify_expressions`                    | Constant-folds and simplifies expressions while preserving output names.                                                    |
+| 5     | `replace_distinct_aggregate`              | Rewrites `DISTINCT` and `DISTINCT ON` operators into aggregate-based plans that later rules can optimize further.           |
+| 6     | `eliminate_join`                          | Replaces keyless inner joins with a literal `false` filter by an empty relation.                                            |
+| 7     | `decorrelate_predicate_subquery`          | Converts eligible `IN` and `EXISTS` predicate subqueries into semi or anti joins.                                           |
+| 8     | `scalar_subquery_to_join`                 | Rewrites eligible scalar subqueries into joins and adds schema-preserving projections.                                      |
+| 9     | `decorrelate_lateral_join`                | Rewrites eligible lateral joins into regular joins.                                                                         |
+| 10    | `extract_equijoin_predicate`              | Splits join filters into equijoin keys and residual predicates.                                                             |
+| 11    | `eliminate_duplicated_expr`               | Removes duplicate expressions from projections, aggregates, and similar operators.                                          |
+| 12    | `eliminate_filter`                        | Drops always-true filters and replaces always-false or NULL filters with empty relations.                                   |
+| 13    | `eliminate_cross_join`                    | Uses filter predicates to replace cross joins with inner joins when join keys can be found.                                 |
+| 14    | `eliminate_limit`                         | Removes no-op limits and simplifies trivial limit shapes.                                                                   |
+| 15    | `propagate_empty_relation`                | Pushes empty-relation knowledge upward so operators fed by no rows collapse early.                                          |
+| 16    | `filter_null_join_keys`                   | Adds `IS NOT NULL` filters to nullable equijoin keys that can never match.                                                  |
+| 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    | `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::Insensitive`] 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::Sensitive
+    }
+
     /// 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, `Insensitive` 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`, ...
+    Insensitive,
+    /// 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.
+    Sensitive,
+    /// 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::Insensitive` — the accumulator
+    // ignores `is_distinct` and returns an unspecified input value either
+    // way. Grouped with `first_value`/`last_value` and left at the default
+    // `Sensitive` 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::{
@@ -870,6 +871,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::Insensitive
+    }
 }
 
 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::Insensitive
+            }
+            // XOR cancels duplicate pairs, so `DISTINCT` is meaningful.
+            BitwiseOperationType::Xor => DistinctHandling::Sensitive,
+        }
+    }
 }
 
 struct BitAndAccumulator<T: ArrowNumericType> {