From ed513cb6ef41ef7356f44274284e5ed34ae40b87 Mon Sep 17 00:00:00 2001 From: Michael Kleen Date: Wed, 16 Sep 2026 19:58:16 +0200 Subject: [PATCH] feat(spark): report distinct_handling for Spark aggregate functions Override `AggregateUDFImpl::distinct_handling` on the Spark aggregates to match what their accumulators actually do: - `collect_set`: `Insensitive`, the accumulator always deduplicates, so `EliminateAggregateDistinct` can drop a redundant `DISTINCT`. - `collect_list`, `try_sum`: `Unsupported`, the accumulators ignore `is_distinct` and silently return the non-distinct result. - `avg`: `Unsupported`, the accumulator rejects `DISTINCT` with an error. Adds unit tests asserting each function's tag and a sqllogictest showing `collect_set(DISTINCT x)` is planned as `collect_set(x)`. Closes #25375 Co-Authored-By: Claude Opus 5 --- .../spark/src/function/aggregate/avg.rs | 17 +++++++-- .../spark/src/function/aggregate/collect.rs | 36 ++++++++++++++++++- .../spark/src/function/aggregate/try_sum.rs | 18 +++++++++- .../test_files/spark/aggregate/collect.slt | 19 ++++++++++ 4 files changed, 86 insertions(+), 4 deletions(-) diff --git a/datafusion/spark/src/function/aggregate/avg.rs b/datafusion/spark/src/function/aggregate/avg.rs index 46e63013dbafb..fc02c0eda4528 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::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 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..aa2fccb37ac21 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,12 @@ impl AggregateUDFImpl for SparkCollectList { fn default_value(&self, data_type: &DataType) -> Result { 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 + } } // @@ -166,6 +174,11 @@ impl AggregateUDFImpl for SparkCollectSet { fn default_value(&self, data_type: &DataType) -> Result { 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 Accumulator for NullToEmptyListAccumulator { 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 + ); + } +} diff --git a/datafusion/spark/src/function/aggregate/try_sum.rs b/datafusion/spark/src/function/aggregate/try_sum.rs index d1f99f4ebc0c3..d73a69f5e896d 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,12 @@ impl AggregateUDFImpl for SparkTrySum { fn default_value(&self, _data_type: &DataType) -> Result { 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(); diff --git a/datafusion/sqllogictest/test_files/spark/aggregate/collect.slt b/datafusion/sqllogictest/test_files/spark/aggregate/collect.slt index 27cab090097da..68f411b756d29 100644 --- a/datafusion/sqllogictest/test_files/spark/aggregate/collect.slt +++ b/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]