diff --git a/datafusion/spark/src/function/aggregate/avg.rs b/datafusion/spark/src/function/aggregate/avg.rs index 38778e49beaf7..4520fb08833b6 100644 --- a/datafusion/spark/src/function/aggregate/avg.rs +++ b/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 Result> { AvgGroupsAccumulator::::new(&DataType::Float64, |sum, count| { Ok(sum / count as f64) diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs index 310bc1c890657..c6c610f5d411f 100644 --- a/datafusion/spark/src/function/aggregate/collect.rs +++ b/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 { empty_list_scalar(data_type) } + + fn distinct_handling(&self) -> DistinctHandling { + DistinctHandling::Unsupported + } } // @@ -166,6 +172,10 @@ impl AggregateUDFImpl for SparkCollectSet { fn default_value(&self, data_type: &DataType) -> Result { 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 Accumulator for NullToEmptyListAccumulator { 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 + ); + } +} diff --git a/datafusion/spark/src/function/aggregate/try_sum.rs b/datafusion/spark/src/function/aggregate/try_sum.rs index d1f99f4ebc0c3..94ff5d7cc9bb9 100644 --- a/datafusion/spark/src/function/aggregate/try_sum.rs +++ b/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 { 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>) -> ArrayRef { Arc::new(Int64Array::from(values)) as ArrayRef } diff --git a/datafusion/sqllogictest/test_files/spark/aggregate/collect.slt b/datafusion/sqllogictest/test_files/spark/aggregate/collect.slt index 27cab090097da..df4380235f356 100644 --- a/datafusion/sqllogictest/test_files/spark/aggregate/collect.slt +++ b/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); ----