fix(spark): report aggregate DISTINCT handling

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::{AggregateOrderSensitivity, 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,
@@ -152,6 +152,10 @@ impl AggregateUDFImpl for SparkAvg {
     fn signature(&self) -> &Signature {
         &self.signature
     }
+
+    fn distinct_handling(&self) -> DistinctHandling {
+        DistinctHandling::Unsupported
+    }
 }
 
 /// An accumulator to compute the average
@@ -381,6 +385,14 @@ mod tests {
     use super::*;
     use arrow::array::Float64Array;
 
+    #[test]
+    fn avg_reports_distinct_as_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,10 @@ impl AggregateUDFImpl for SparkCollectList {
     fn default_value(&self, data_type: &DataType) -> Result<ScalarValue> {
         empty_list_scalar(data_type)
     }
+
+    fn distinct_handling(&self) -> DistinctHandling {
+        DistinctHandling::Unsupported
+    }
 }
 
 // <https://spark.apache.org/docs/latest/api/sql/index.html#collect_set>
@@ -166,6 +172,10 @@ impl AggregateUDFImpl for SparkCollectSet {
     fn default_value(&self, data_type: &DataType) -> Result<ScalarValue> {
         empty_list_scalar(data_type)
     }
+
+    fn distinct_handling(&self) -> DistinctHandling {
+        DistinctHandling::Insensitive
+    }
 }
 
 /// Wrapper accumulator that returns an empty list instead of NULL when all inputs are NULL.
@@ -216,3 +226,20 @@ impl<T: Accumulator> Accumulator for NullToEmptyListAccumulator<T> {
         self.inner.size() + self.list_type.size()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn collect_aggregates_report_distinct_handling() {
+        assert_eq!(
+            SparkCollectList::new().distinct_handling(),
+            DistinctHandling::Unsupported
+        );
+        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,10 @@ impl AggregateUDFImpl for SparkTrySum {
     fn default_value(&self, _data_type: &DataType) -> Result<ScalarValue> {
         Ok(ScalarValue::Null)
     }
+
+    fn distinct_handling(&self) -> DistinctHandling {
+        DistinctHandling::Unsupported
+    }
 }
 
 #[cfg(test)]
@@ -341,6 +347,14 @@ mod tests {
     use super::*;
     // -------- Helpers --------
 
+    #[test]
+    fn try_sum_reports_distinct_as_unsupported() {
+        assert_eq!(
+            SparkTrySum::new().distinct_handling(),
+            DistinctHandling::Unsupported
+        );
+    }
+
     fn int64(values: Vec<Option<i64>>) -> ArrayRef {
         Arc::new(Int64Array::from(values)) as ArrayRef
     }

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

@@ -64,6 +64,24 @@ SELECT array_sort(collect_set(a)) FROM (VALUES (1), (2), (2), (3), (1)) AS t(a);
 ----
 [1, 2, 3]
 
+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)) 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))
+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]
+
 query ?
 SELECT array_sort(collect_set(a)) FROM (VALUES (1), (NULL), (3)) AS t(a);
 ----