Skip to content

Commit ed513cb

Browse files
mkleenclaude
andcommitted
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 <noreply@anthropic.com>
1 parent 140c7c5 commit ed513cb

4 files changed

Lines changed: 86 additions & 4 deletions

File tree

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

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,8 @@ use datafusion_common::{Result, ScalarValue, not_impl_err};
2929
use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
3030
use datafusion_expr::utils::format_state_name;
3131
use datafusion_expr::{
32-
Accumulator, AggregateUDFImpl, Coercion, EmitTo, GroupsAccumulator, ReversedUDAF,
33-
Signature, TypeSignatureClass, Volatility,
32+
Accumulator, AggregateUDFImpl, Coercion, DistinctHandling, EmitTo, GroupsAccumulator,
33+
ReversedUDAF, Signature, TypeSignatureClass, Volatility,
3434
};
3535
use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::{
3636
filtered_null_mask, set_nulls,
@@ -148,6 +148,11 @@ impl AggregateUDFImpl for SparkAvg {
148148
fn signature(&self) -> &Signature {
149149
&self.signature
150150
}
151+
fn distinct_handling(&self) -> DistinctHandling {
152+
// Duplicate-sensitive, and the accumulator rejects `DISTINCT` with
153+
// a not-implemented error.
154+
DistinctHandling::Unsupported
155+
}
151156
}
152157

153158
/// An accumulator to compute the average
@@ -377,6 +382,14 @@ mod tests {
377382
use super::*;
378383
use arrow::array::Float64Array;
379384

385+
#[test]
386+
fn distinct_handling_is_unsupported() {
387+
assert_eq!(
388+
SparkAvg::new().distinct_handling(),
389+
DistinctHandling::Unsupported
390+
);
391+
}
392+
380393
fn make_acc() -> AvgGroupsAccumulator<Float64Type, impl Fn(f64, i64) -> Result<f64>> {
381394
AvgGroupsAccumulator::<Float64Type, _>::new(&DataType::Float64, |sum, count| {
382395
Ok(sum / count as f64)

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

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,9 @@ use datafusion_common::utils::SingleRowListArrayBuilder;
2121
use datafusion_common::{Result, ScalarValue, internal_err};
2222
use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
2323
use datafusion_expr::utils::format_state_name;
24-
use datafusion_expr::{Accumulator, AggregateUDFImpl, Signature, Volatility};
24+
use datafusion_expr::{
25+
Accumulator, AggregateUDFImpl, DistinctHandling, Signature, Volatility,
26+
};
2527
use datafusion_functions_aggregate::array_agg::{
2628
ArrayAggAccumulator, DistinctArrayAggAccumulator,
2729
};
@@ -105,6 +107,12 @@ impl AggregateUDFImpl for SparkCollectList {
105107
fn default_value(&self, data_type: &DataType) -> Result<ScalarValue> {
106108
empty_list_scalar(data_type)
107109
}
110+
fn distinct_handling(&self) -> DistinctHandling {
111+
// Duplicate-sensitive, but the accumulator does not read
112+
// `is_distinct` and today silently returns the non-distinct answer.
113+
// The tag records the intent; enforcement is a follow-up change.
114+
DistinctHandling::Unsupported
115+
}
108116
}
109117

110118
// <https://spark.apache.org/docs/latest/api/sql/index.html#collect_set>
@@ -166,6 +174,11 @@ impl AggregateUDFImpl for SparkCollectSet {
166174
fn default_value(&self, data_type: &DataType) -> Result<ScalarValue> {
167175
empty_list_scalar(data_type)
168176
}
177+
fn distinct_handling(&self) -> DistinctHandling {
178+
// The accumulator always deduplicates, so `DISTINCT` cannot change
179+
// the result.
180+
DistinctHandling::Insensitive
181+
}
169182
}
170183

171184
/// 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> {
216229
self.inner.size() + self.list_type.size()
217230
}
218231
}
232+
233+
#[cfg(test)]
234+
mod tests {
235+
use super::*;
236+
237+
#[test]
238+
fn collect_list_distinct_handling() {
239+
assert_eq!(
240+
SparkCollectList::new().distinct_handling(),
241+
DistinctHandling::Unsupported
242+
);
243+
}
244+
245+
#[test]
246+
fn collect_set_distinct_handling() {
247+
assert_eq!(
248+
SparkCollectSet::new().distinct_handling(),
249+
DistinctHandling::Insensitive
250+
);
251+
}
252+
}

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

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,9 @@ use arrow::datatypes::{
2323
use datafusion_common::{Result, ScalarValue, downcast_value, exec_err, not_impl_err};
2424
use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
2525
use datafusion_expr::utils::format_state_name;
26-
use datafusion_expr::{Accumulator, AggregateUDFImpl, Signature, Volatility};
26+
use datafusion_expr::{
27+
Accumulator, AggregateUDFImpl, DistinctHandling, Signature, Volatility,
28+
};
2729
use std::fmt::{Debug, Formatter};
2830
use std::mem::size_of_val;
2931

@@ -330,6 +332,12 @@ impl AggregateUDFImpl for SparkTrySum {
330332
fn default_value(&self, _data_type: &DataType) -> Result<ScalarValue> {
331333
Ok(ScalarValue::Null)
332334
}
335+
fn distinct_handling(&self) -> DistinctHandling {
336+
// Duplicate-sensitive, but the accumulator does not read
337+
// `is_distinct` and today silently returns the non-distinct answer.
338+
// The tag records the intent; enforcement is a follow-up change.
339+
DistinctHandling::Unsupported
340+
}
333341
}
334342

335343
#[cfg(test)]
@@ -632,6 +640,14 @@ mod tests {
632640
Ok(())
633641
}
634642

643+
#[test]
644+
fn distinct_handling_is_unsupported() {
645+
assert_eq!(
646+
SparkTrySum::new().distinct_handling(),
647+
DistinctHandling::Unsupported
648+
);
649+
}
650+
635651
#[test]
636652
fn decimal_38_0_max_precision_overflows_to_null() -> Result<()> {
637653
let f = SparkTrySum::new();

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,3 +122,22 @@ ORDER BY g;
122122
1 [5]
123123
2 [10, 20]
124124
3 []
125+
126+
# collect_set always deduplicates, so DISTINCT is redundant and removed by the optimizer
127+
query ?
128+
SELECT array_sort(collect_set(DISTINCT a)) FROM (VALUES (1), (2), (2), (3), (1)) AS t(a);
129+
----
130+
[1, 2, 3]
131+
132+
query TT
133+
EXPLAIN SELECT collect_set(DISTINCT a) FROM (VALUES (1), (2), (2), (3), (1)) AS t(a);
134+
----
135+
logical_plan
136+
01)Aggregate: groupBy=[[]], aggr=[[collect_set(t.a) AS collect_set(DISTINCT t.a)]]
137+
02)--SubqueryAlias: t
138+
03)----Projection: column1 AS a
139+
04)------Values: (Int64(1)), (Int64(2)), (Int64(2)), (Int64(3)), (Int64(1))
140+
physical_plan
141+
01)AggregateExec: mode=Single, gby=[], aggr=[collect_set(t.a) as collect_set(DISTINCT t.a)]
142+
02)--ProjectionExec: expr=[column1@0 as a]
143+
03)----DataSourceExec: partitions=1, partition_sizes=[1]

0 commit comments

Comments
 (0)