datafusion_spark/function/aggregate/
collect.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use arrow::array::ArrayRef;
19use arrow::datatypes::{DataType, Field, FieldRef};
20use datafusion_common::utils::SingleRowListArrayBuilder;
21use datafusion_common::{Result, ScalarValue, internal_err};
22use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
23use datafusion_expr::utils::format_state_name;
24use datafusion_expr::{Accumulator, AggregateUDFImpl, Signature, Volatility};
25use datafusion_functions_aggregate::array_agg::{
26    ArrayAggAccumulator, DistinctArrayAggAccumulator,
27};
28use std::sync::Arc;
29
30// Spark implementation of collect_list/collect_set aggregate function.
31// Differs from DataFusion ArrayAgg in the following ways:
32// - ignores NULL inputs
33// - returns an empty list when all inputs are NULL
34// - does not support ordering
35
36/// Build an empty list `ScalarValue` for a `List(element_type)` data type.
37/// Used as the result for empty window frames and for groups whose inputs
38/// were all NULL, matching Spark's `collect_list` / `collect_set` semantics.
39fn empty_list_scalar(list_type: &DataType) -> Result<ScalarValue> {
40    let DataType::List(field) = list_type else {
41        return internal_err!(
42            "collect_list/collect_set expected List return type, got {list_type:?}"
43        );
44    };
45    let empty = arrow::array::new_empty_array(field.data_type());
46    Ok(SingleRowListArrayBuilder::new(empty).build_list_scalar())
47}
48
49// <https://spark.apache.org/docs/latest/api/sql/index.html#collect_list>
50#[derive(Debug, PartialEq, Eq, Hash)]
51pub struct SparkCollectList {
52    signature: Signature,
53}
54
55impl Default for SparkCollectList {
56    fn default() -> Self {
57        Self::new()
58    }
59}
60
61impl SparkCollectList {
62    pub fn new() -> Self {
63        Self {
64            signature: Signature::any(1, Volatility::Immutable),
65        }
66    }
67}
68
69impl AggregateUDFImpl for SparkCollectList {
70    fn name(&self) -> &str {
71        "collect_list"
72    }
73
74    fn signature(&self) -> &Signature {
75        &self.signature
76    }
77
78    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
79        Ok(DataType::List(Arc::new(Field::new_list_field(
80            arg_types[0].clone(),
81            true,
82        ))))
83    }
84
85    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
86        Ok(vec![
87            Field::new_list(
88                format_state_name(args.name, "collect_list"),
89                Field::new_list_field(args.input_fields[0].data_type().clone(), true),
90                true,
91            )
92            .into(),
93        ])
94    }
95
96    fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
97        let element_type = acc_args.expr_fields[0].data_type().clone();
98        let ignore_nulls = true;
99        Ok(Box::new(NullToEmptyListAccumulator::new(
100            ArrayAggAccumulator::try_new(&element_type, ignore_nulls)?,
101            acc_args.return_type().clone(),
102        )))
103    }
104
105    fn default_value(&self, data_type: &DataType) -> Result<ScalarValue> {
106        empty_list_scalar(data_type)
107    }
108}
109
110// <https://spark.apache.org/docs/latest/api/sql/index.html#collect_set>
111#[derive(Debug, PartialEq, Eq, Hash)]
112pub struct SparkCollectSet {
113    signature: Signature,
114}
115
116impl Default for SparkCollectSet {
117    fn default() -> Self {
118        Self::new()
119    }
120}
121
122impl SparkCollectSet {
123    pub fn new() -> Self {
124        Self {
125            signature: Signature::any(1, Volatility::Immutable),
126        }
127    }
128}
129
130impl AggregateUDFImpl for SparkCollectSet {
131    fn name(&self) -> &str {
132        "collect_set"
133    }
134
135    fn signature(&self) -> &Signature {
136        &self.signature
137    }
138
139    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
140        Ok(DataType::List(Arc::new(Field::new_list_field(
141            arg_types[0].clone(),
142            true,
143        ))))
144    }
145
146    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
147        Ok(vec![
148            Field::new_list(
149                format_state_name(args.name, "collect_set"),
150                Field::new_list_field(args.input_fields[0].data_type().clone(), true),
151                true,
152            )
153            .into(),
154        ])
155    }
156
157    fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
158        let element_type = acc_args.expr_fields[0].data_type().clone();
159        let ignore_nulls = true;
160        Ok(Box::new(NullToEmptyListAccumulator::new(
161            DistinctArrayAggAccumulator::try_new(&element_type, None, ignore_nulls)?,
162            acc_args.return_type().clone(),
163        )))
164    }
165
166    fn default_value(&self, data_type: &DataType) -> Result<ScalarValue> {
167        empty_list_scalar(data_type)
168    }
169}
170
171/// Wrapper accumulator that returns an empty list instead of NULL when all inputs are NULL.
172/// This implements Spark's behavior for collect_list and collect_set.
173#[derive(Debug)]
174struct NullToEmptyListAccumulator<T: Accumulator> {
175    inner: T,
176    list_type: DataType,
177}
178
179impl<T: Accumulator> NullToEmptyListAccumulator<T> {
180    pub fn new(inner: T, list_type: DataType) -> Self {
181        Self { inner, list_type }
182    }
183}
184
185impl<T: Accumulator> Accumulator for NullToEmptyListAccumulator<T> {
186    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
187        self.inner.update_batch(values)
188    }
189
190    fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
191        self.inner.merge_batch(states)
192    }
193
194    fn state(&mut self) -> Result<Vec<ScalarValue>> {
195        self.inner.state()
196    }
197
198    fn evaluate(&mut self) -> Result<ScalarValue> {
199        let result = self.inner.evaluate()?;
200        if result.is_null() {
201            empty_list_scalar(&self.list_type)
202        } else {
203            Ok(result)
204        }
205    }
206
207    fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
208        self.inner.retract_batch(values)
209    }
210
211    fn supports_retract_batch(&self) -> bool {
212        self.inner.supports_retract_batch()
213    }
214
215    fn size(&self) -> usize {
216        self.inner.size() + self.list_type.size()
217    }
218}
datafusion_spark/function/aggregate/collect.rs

datafusion_spark/function/aggregate/
collect.rs
