feat(spark): add distinct handling for Spark aggregate functions

datafusion/spark/src/function/aggregate/avg.rs

@@ -29,8 +29,8 @@ use datafusion_common::{Result, ScalarValue, not_impl_err};
 use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
 use datafusion_expr::utils::format_state_name;
 use datafusion_expr::{
-    Accumulator, AggregateUDFImpl, Coercion, EmitTo, GroupsAccumulator, ReversedUDAF,
-    Signature, TypeSignatureClass, Volatility,
+    Accumulator, AggregateUDFImpl, Coercion, DistinctHandling, EmitTo, GroupsAccumulator,
+    ReversedUDAF, Signature, TypeSignatureClass, Volatility,
 };
 use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::{
     filtered_null_mask, set_nulls,
@@ -148,6 +148,11 @@ impl AggregateUDFImpl for SparkAvg {
     fn signature(&self) -> &Signature {
         &self.signature
     }
+    fn distinct_handling(&self) -> DistinctHandling {
+        // Duplicate-sensitive, and the accumulator rejects `DISTINCT` with
+        // a not-implemented error.
+        DistinctHandling::Unsupported
+    }
 }
 
 /// An accumulator to compute the average
@@ -377,6 +382,14 @@ mod tests {
     use super::*;
     use arrow::array::Float64Array;
 
+    #[test]
+    fn distinct_handling_is_unsupported() {
+        assert_eq!(
+            SparkAvg::new().distinct_handling(),
+            DistinctHandling::Unsupported
+        );
+    }
+
     fn make_acc() -> AvgGroupsAccumulator<Float64Type, impl Fn(f64, i64) -> Result<f64>> {
         AvgGroupsAccumulator::<Float64Type, _>::new(&DataType::Float64, |sum, count| {
             Ok(sum / count as f64)

datafusion/spark/src/function/aggregate/collect.rs

@@ -21,7 +21,9 @@ use datafusion_common::utils::SingleRowListArrayBuilder;
 use datafusion_common::{Result, ScalarValue, internal_err};
 use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
 use datafusion_expr::utils::format_state_name;
-use datafusion_expr::{Accumulator, AggregateUDFImpl, Signature, Volatility};
+use datafusion_expr::{
+    Accumulator, AggregateUDFImpl, DistinctHandling, Signature, Volatility,
+};
 use datafusion_functions_aggregate::array_agg::{
     ArrayAggAccumulator, DistinctArrayAggAccumulator,
 };
@@ -105,6 +107,12 @@ impl AggregateUDFImpl for SparkCollectList {
     fn default_value(&self, data_type: &DataType) -> Result<ScalarValue> {
         empty_list_scalar(data_type)
     }
+    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
+    }
 }
 
 // <https://spark.apache.org/docs/latest/api/sql/index.html#collect_set>
@@ -166,6 +174,11 @@ impl AggregateUDFImpl for SparkCollectSet {
     fn default_value(&self, data_type: &DataType) -> Result<ScalarValue> {
         empty_list_scalar(data_type)
     }
+    fn distinct_handling(&self) -> DistinctHandling {
+        // The accumulator always deduplicates, so `DISTINCT` cannot change
+        // the result.
+        DistinctHandling::Insensitive
+    }
 }
 
 /// Wrapper accumulator that returns an empty list instead of NULL when all inputs are NULL.
@@ -216,3 +229,24 @@ impl<T: Accumulator> Accumulator for NullToEmptyListAccumulator<T> {
         self.inner.size() + self.list_type.size()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn collect_list_distinct_handling() {
+        assert_eq!(
+            SparkCollectList::new().distinct_handling(),
+            DistinctHandling::Unsupported
+        );
+    }
+
+    #[test]
+    fn collect_set_distinct_handling() {
+        assert_eq!(
+            SparkCollectSet::new().distinct_handling(),
+            DistinctHandling::Insensitive
+        );
+    }
+}

datafusion/spark/src/function/aggregate/try_sum.rs

@@ -23,7 +23,9 @@ use arrow::datatypes::{
 use datafusion_common::{Result, ScalarValue, downcast_value, exec_err, not_impl_err};
 use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
 use datafusion_expr::utils::format_state_name;
-use datafusion_expr::{Accumulator, AggregateUDFImpl, Signature, Volatility};
+use datafusion_expr::{
+    Accumulator, AggregateUDFImpl, DistinctHandling, Signature, Volatility,
+};
 use std::fmt::{Debug, Formatter};
 use std::mem::size_of_val;
 
@@ -330,6 +332,12 @@ impl AggregateUDFImpl for SparkTrySum {
     fn default_value(&self, _data_type: &DataType) -> Result<ScalarValue> {
         Ok(ScalarValue::Null)
     }
+    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
+    }
 }
 
 #[cfg(test)]
@@ -632,6 +640,14 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn distinct_handling_is_unsupported() {
+        assert_eq!(
+            SparkTrySum::new().distinct_handling(),
+            DistinctHandling::Unsupported
+        );
+    }
+
     #[test]
     fn decimal_38_0_max_precision_overflows_to_null() -> Result<()> {
         let f = SparkTrySum::new();

datafusion/sqllogictest/test_files/spark/aggregate/collect.slt

@@ -122,3 +122,22 @@ ORDER BY g;
 1 [5]
 2 [10, 20]
 3 []
+
+# collect_set always deduplicates, so DISTINCT is redundant and removed by the optimizer
+query ?
+SELECT array_sort(collect_set(DISTINCT a)) FROM (VALUES (1), (2), (2), (3), (1)) AS t(a);
+----
+[1, 2, 3]
+
+query TT
+EXPLAIN SELECT collect_set(DISTINCT a) FROM (VALUES (1), (2), (2), (3), (1)) AS t(a);
+----
+logical_plan
+01)Aggregate: groupBy=[[]], aggr=[[collect_set(t.a) AS collect_set(DISTINCT t.a)]]
+02)--SubqueryAlias: t
+03)----Projection: column1 AS a
+04)------Values: (Int64(1)), (Int64(2)), (Int64(2)), (Int64(3)), (Int64(1))
+physical_plan
+01)AggregateExec: mode=Single, gby=[], aggr=[collect_set(t.a) as collect_set(DISTINCT t.a)]
+02)--ProjectionExec: expr=[column1@0 as a]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]