datafusion_functions_aggregate/
string_agg.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
18//! [`StringAgg`] accumulator for the `string_agg` function
19
20use std::any::Any;
21use std::hash::Hash;
22use std::mem::size_of_val;
23
24use crate::array_agg::ArrayAgg;
25
26use arrow::array::ArrayRef;
27use arrow::datatypes::{DataType, Field, FieldRef};
28use datafusion_common::cast::{as_generic_string_array, as_string_view_array};
29use datafusion_common::{internal_err, not_impl_err, Result, ScalarValue};
30use datafusion_expr::function::AccumulatorArgs;
31use datafusion_expr::{
32    Accumulator, AggregateUDFImpl, Documentation, Signature, TypeSignature, Volatility,
33};
34use datafusion_functions_aggregate_common::accumulator::StateFieldsArgs;
35use datafusion_macros::user_doc;
36use datafusion_physical_expr::expressions::Literal;
37
38make_udaf_expr_and_func!(
39    StringAgg,
40    string_agg,
41    expr delimiter,
42    "Concatenates the values of string expressions and places separator values between them",
43    string_agg_udaf
44);
45
46#[user_doc(
47    doc_section(label = "General Functions"),
48    description = "Concatenates the values of string expressions and places separator values between them. \
49If ordering is required, strings are concatenated in the specified order. \
50This aggregation function can only mix DISTINCT and ORDER BY if the ordering expression is exactly the same as the first argument expression.",
51    syntax_example = "string_agg([DISTINCT] expression, delimiter [ORDER BY expression])",
52    sql_example = r#"```sql
53> SELECT string_agg(name, ', ') AS names_list
54  FROM employee;
55+--------------------------+
56| names_list               |
57+--------------------------+
58| Alice, Bob, Bob, Charlie |
59+--------------------------+
60> SELECT string_agg(name, ', ' ORDER BY name DESC) AS names_list
61  FROM employee;
62+--------------------------+
63| names_list               |
64+--------------------------+
65| Charlie, Bob, Bob, Alice |
66+--------------------------+
67> SELECT string_agg(DISTINCT name, ', ' ORDER BY name DESC) AS names_list
68  FROM employee;
69+--------------------------+
70| names_list               |
71+--------------------------+
72| Charlie, Bob, Alice |
73+--------------------------+
74```"#,
75    argument(
76        name = "expression",
77        description = "The string expression to concatenate. Can be a column or any valid string expression."
78    ),
79    argument(
80        name = "delimiter",
81        description = "A literal string used as a separator between the concatenated values."
82    )
83)]
84/// STRING_AGG aggregate expression
85#[derive(Debug, PartialEq, Eq, Hash)]
86pub struct StringAgg {
87    signature: Signature,
88    array_agg: ArrayAgg,
89}
90
91impl StringAgg {
92    /// Create a new StringAgg aggregate function
93    pub fn new() -> Self {
94        Self {
95            signature: Signature::one_of(
96                vec![
97                    TypeSignature::Exact(vec![DataType::LargeUtf8, DataType::Utf8]),
98                    TypeSignature::Exact(vec![DataType::LargeUtf8, DataType::LargeUtf8]),
99                    TypeSignature::Exact(vec![DataType::LargeUtf8, DataType::Null]),
100                    TypeSignature::Exact(vec![DataType::LargeUtf8, DataType::Utf8View]),
101                    TypeSignature::Exact(vec![DataType::Utf8, DataType::Utf8]),
102                    TypeSignature::Exact(vec![DataType::Utf8, DataType::LargeUtf8]),
103                    TypeSignature::Exact(vec![DataType::Utf8, DataType::Null]),
104                    TypeSignature::Exact(vec![DataType::Utf8, DataType::Utf8View]),
105                    TypeSignature::Exact(vec![DataType::Utf8View, DataType::Utf8View]),
106                    TypeSignature::Exact(vec![DataType::Utf8View, DataType::LargeUtf8]),
107                    TypeSignature::Exact(vec![DataType::Utf8View, DataType::Null]),
108                    TypeSignature::Exact(vec![DataType::Utf8View, DataType::Utf8]),
109                ],
110                Volatility::Immutable,
111            ),
112            array_agg: Default::default(),
113        }
114    }
115}
116
117impl Default for StringAgg {
118    fn default() -> Self {
119        Self::new()
120    }
121}
122
123impl AggregateUDFImpl for StringAgg {
124    fn as_any(&self) -> &dyn Any {
125        self
126    }
127
128    fn name(&self) -> &str {
129        "string_agg"
130    }
131
132    fn signature(&self) -> &Signature {
133        &self.signature
134    }
135
136    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
137        Ok(DataType::LargeUtf8)
138    }
139
140    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
141        self.array_agg.state_fields(args)
142    }
143
144    fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
145        let Some(lit) = acc_args.exprs[1].as_any().downcast_ref::<Literal>() else {
146            return not_impl_err!(
147                "The second argument of the string_agg function must be a string literal"
148            );
149        };
150
151        let delimiter = if lit.value().is_null() {
152            // If the second argument (the delimiter that joins strings) is NULL, join
153            // on an empty string. (e.g. [a, b, c] => "abc").
154            ""
155        } else if let Some(lit_string) = lit.value().try_as_str() {
156            lit_string.unwrap_or("")
157        } else {
158            return not_impl_err!(
159                "StringAgg not supported for delimiter \"{}\"",
160                lit.value()
161            );
162        };
163
164        let array_agg_acc = self.array_agg.accumulator(AccumulatorArgs {
165            return_field: Field::new(
166                "f",
167                DataType::new_list(acc_args.return_field.data_type().clone(), true),
168                true,
169            )
170            .into(),
171            exprs: &filter_index(acc_args.exprs, 1),
172            ..acc_args
173        })?;
174
175        Ok(Box::new(StringAggAccumulator::new(
176            array_agg_acc,
177            delimiter,
178        )))
179    }
180
181    fn reverse_expr(&self) -> datafusion_expr::ReversedUDAF {
182        datafusion_expr::ReversedUDAF::Reversed(string_agg_udaf())
183    }
184
185    fn documentation(&self) -> Option<&Documentation> {
186        self.doc()
187    }
188}
189
190#[derive(Debug)]
191pub(crate) struct StringAggAccumulator {
192    array_agg_acc: Box<dyn Accumulator>,
193    delimiter: String,
194}
195
196impl StringAggAccumulator {
197    pub fn new(array_agg_acc: Box<dyn Accumulator>, delimiter: &str) -> Self {
198        Self {
199            array_agg_acc,
200            delimiter: delimiter.to_string(),
201        }
202    }
203}
204
205impl Accumulator for StringAggAccumulator {
206    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
207        self.array_agg_acc.update_batch(&filter_index(values, 1))
208    }
209
210    fn evaluate(&mut self) -> Result<ScalarValue> {
211        let scalar = self.array_agg_acc.evaluate()?;
212
213        let ScalarValue::List(list) = scalar else {
214            return internal_err!("Expected a DataType::List while evaluating underlying ArrayAggAccumulator, but got {}", scalar.data_type());
215        };
216
217        let string_arr: Vec<_> = match list.value_type() {
218            DataType::LargeUtf8 => as_generic_string_array::<i64>(list.values())?
219                .iter()
220                .flatten()
221                .collect(),
222            DataType::Utf8 => as_generic_string_array::<i32>(list.values())?
223                .iter()
224                .flatten()
225                .collect(),
226            DataType::Utf8View => as_string_view_array(list.values())?
227                .iter()
228                .flatten()
229                .collect(),
230            _ => {
231                return internal_err!(
232                    "Expected elements to of type Utf8 or LargeUtf8, but got {}",
233                    list.value_type()
234                )
235            }
236        };
237
238        if string_arr.is_empty() {
239            return Ok(ScalarValue::LargeUtf8(None));
240        }
241
242        Ok(ScalarValue::LargeUtf8(Some(
243            string_arr.join(&self.delimiter),
244        )))
245    }
246
247    fn size(&self) -> usize {
248        size_of_val(self) - size_of_val(&self.array_agg_acc)
249            + self.array_agg_acc.size()
250            + self.delimiter.capacity()
251    }
252
253    fn state(&mut self) -> Result<Vec<ScalarValue>> {
254        self.array_agg_acc.state()
255    }
256
257    fn merge_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
258        self.array_agg_acc.merge_batch(values)
259    }
260}
261
262fn filter_index<T: Clone>(values: &[T], index: usize) -> Vec<T> {
263    values
264        .iter()
265        .enumerate()
266        .filter(|(i, _)| *i != index)
267        .map(|(_, v)| v)
268        .cloned()
269        .collect::<Vec<_>>()
270}
271
272#[cfg(test)]
273mod tests {
274    use super::*;
275    use arrow::array::LargeStringArray;
276    use arrow::compute::SortOptions;
277    use arrow::datatypes::{Fields, Schema};
278    use datafusion_common::internal_err;
279    use datafusion_physical_expr::expressions::Column;
280    use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
281    use std::sync::Arc;
282
283    #[test]
284    fn no_duplicates_no_distinct() -> Result<()> {
285        let (mut acc1, mut acc2) = StringAggAccumulatorBuilder::new(",").build_two()?;
286
287        acc1.update_batch(&[data(["a", "b", "c"]), data([","])])?;
288        acc2.update_batch(&[data(["d", "e", "f"]), data([","])])?;
289        acc1 = merge(acc1, acc2)?;
290
291        let result = some_str(acc1.evaluate()?);
292
293        assert_eq!(result, "a,b,c,d,e,f");
294
295        Ok(())
296    }
297
298    #[test]
299    fn no_duplicates_distinct() -> Result<()> {
300        let (mut acc1, mut acc2) = StringAggAccumulatorBuilder::new(",")
301            .distinct()
302            .build_two()?;
303
304        acc1.update_batch(&[data(["a", "b", "c"]), data([","])])?;
305        acc2.update_batch(&[data(["d", "e", "f"]), data([","])])?;
306        acc1 = merge(acc1, acc2)?;
307
308        let result = some_str_sorted(acc1.evaluate()?, ",");
309
310        assert_eq!(result, "a,b,c,d,e,f");
311
312        Ok(())
313    }
314
315    #[test]
316    fn duplicates_no_distinct() -> Result<()> {
317        let (mut acc1, mut acc2) = StringAggAccumulatorBuilder::new(",").build_two()?;
318
319        acc1.update_batch(&[data(["a", "b", "c"]), data([","])])?;
320        acc2.update_batch(&[data(["a", "b", "c"]), data([","])])?;
321        acc1 = merge(acc1, acc2)?;
322
323        let result = some_str(acc1.evaluate()?);
324
325        assert_eq!(result, "a,b,c,a,b,c");
326
327        Ok(())
328    }
329
330    #[test]
331    fn duplicates_distinct() -> Result<()> {
332        let (mut acc1, mut acc2) = StringAggAccumulatorBuilder::new(",")
333            .distinct()
334            .build_two()?;
335
336        acc1.update_batch(&[data(["a", "b", "c"]), data([","])])?;
337        acc2.update_batch(&[data(["a", "b", "c"]), data([","])])?;
338        acc1 = merge(acc1, acc2)?;
339
340        let result = some_str_sorted(acc1.evaluate()?, ",");
341
342        assert_eq!(result, "a,b,c");
343
344        Ok(())
345    }
346
347    #[test]
348    fn no_duplicates_distinct_sort_asc() -> Result<()> {
349        let (mut acc1, mut acc2) = StringAggAccumulatorBuilder::new(",")
350            .distinct()
351            .order_by_col("col", SortOptions::new(false, false))
352            .build_two()?;
353
354        acc1.update_batch(&[data(["e", "b", "d"]), data([","])])?;
355        acc2.update_batch(&[data(["f", "a", "c"]), data([","])])?;
356        acc1 = merge(acc1, acc2)?;
357
358        let result = some_str(acc1.evaluate()?);
359
360        assert_eq!(result, "a,b,c,d,e,f");
361
362        Ok(())
363    }
364
365    #[test]
366    fn no_duplicates_distinct_sort_desc() -> Result<()> {
367        let (mut acc1, mut acc2) = StringAggAccumulatorBuilder::new(",")
368            .distinct()
369            .order_by_col("col", SortOptions::new(true, false))
370            .build_two()?;
371
372        acc1.update_batch(&[data(["e", "b", "d"]), data([","])])?;
373        acc2.update_batch(&[data(["f", "a", "c"]), data([","])])?;
374        acc1 = merge(acc1, acc2)?;
375
376        let result = some_str(acc1.evaluate()?);
377
378        assert_eq!(result, "f,e,d,c,b,a");
379
380        Ok(())
381    }
382
383    #[test]
384    fn duplicates_distinct_sort_asc() -> Result<()> {
385        let (mut acc1, mut acc2) = StringAggAccumulatorBuilder::new(",")
386            .distinct()
387            .order_by_col("col", SortOptions::new(false, false))
388            .build_two()?;
389
390        acc1.update_batch(&[data(["a", "c", "b"]), data([","])])?;
391        acc2.update_batch(&[data(["b", "c", "a"]), data([","])])?;
392        acc1 = merge(acc1, acc2)?;
393
394        let result = some_str(acc1.evaluate()?);
395
396        assert_eq!(result, "a,b,c");
397
398        Ok(())
399    }
400
401    #[test]
402    fn duplicates_distinct_sort_desc() -> Result<()> {
403        let (mut acc1, mut acc2) = StringAggAccumulatorBuilder::new(",")
404            .distinct()
405            .order_by_col("col", SortOptions::new(true, false))
406            .build_two()?;
407
408        acc1.update_batch(&[data(["a", "c", "b"]), data([","])])?;
409        acc2.update_batch(&[data(["b", "c", "a"]), data([","])])?;
410        acc1 = merge(acc1, acc2)?;
411
412        let result = some_str(acc1.evaluate()?);
413
414        assert_eq!(result, "c,b,a");
415
416        Ok(())
417    }
418
419    struct StringAggAccumulatorBuilder {
420        sep: String,
421        distinct: bool,
422        order_bys: Vec<PhysicalSortExpr>,
423        schema: Schema,
424    }
425
426    impl StringAggAccumulatorBuilder {
427        fn new(sep: &str) -> Self {
428            Self {
429                sep: sep.to_string(),
430                distinct: Default::default(),
431                order_bys: vec![],
432                schema: Schema {
433                    fields: Fields::from(vec![Field::new(
434                        "col",
435                        DataType::LargeUtf8,
436                        true,
437                    )]),
438                    metadata: Default::default(),
439                },
440            }
441        }
442        fn distinct(mut self) -> Self {
443            self.distinct = true;
444            self
445        }
446
447        fn order_by_col(mut self, col: &str, sort_options: SortOptions) -> Self {
448            self.order_bys.extend([PhysicalSortExpr::new(
449                Arc::new(
450                    Column::new_with_schema(col, &self.schema)
451                        .expect("column not available in schema"),
452                ),
453                sort_options,
454            )]);
455            self
456        }
457
458        fn build(&self) -> Result<Box<dyn Accumulator>> {
459            StringAgg::new().accumulator(AccumulatorArgs {
460                return_field: Field::new("f", DataType::LargeUtf8, true).into(),
461                schema: &self.schema,
462                ignore_nulls: false,
463                order_bys: &self.order_bys,
464                is_reversed: false,
465                name: "",
466                is_distinct: self.distinct,
467                exprs: &[
468                    Arc::new(Column::new("col", 0)),
469                    Arc::new(Literal::new(ScalarValue::Utf8(Some(self.sep.to_string())))),
470                ],
471            })
472        }
473
474        fn build_two(&self) -> Result<(Box<dyn Accumulator>, Box<dyn Accumulator>)> {
475            Ok((self.build()?, self.build()?))
476        }
477    }
478
479    fn some_str(value: ScalarValue) -> String {
480        str(value)
481            .expect("ScalarValue was not a String")
482            .expect("ScalarValue was None")
483    }
484
485    fn some_str_sorted(value: ScalarValue, sep: &str) -> String {
486        let value = some_str(value);
487        let mut parts: Vec<&str> = value.split(sep).collect();
488        parts.sort();
489        parts.join(sep)
490    }
491
492    fn str(value: ScalarValue) -> Result<Option<String>> {
493        match value {
494            ScalarValue::LargeUtf8(v) => Ok(v),
495            _ => internal_err!(
496                "Expected ScalarValue::LargeUtf8, got {}",
497                value.data_type()
498            ),
499        }
500    }
501
502    fn data<const N: usize>(list: [&str; N]) -> ArrayRef {
503        Arc::new(LargeStringArray::from(list.to_vec()))
504    }
505
506    fn merge(
507        mut acc1: Box<dyn Accumulator>,
508        mut acc2: Box<dyn Accumulator>,
509    ) -> Result<Box<dyn Accumulator>> {
510        let intermediate_state = acc2.state().and_then(|e| {
511            e.iter()
512                .map(|v| v.to_array())
513                .collect::<Result<Vec<ArrayRef>>>()
514        })?;
515        acc1.merge_batch(&intermediate_state)?;
516        Ok(acc1)
517    }
518}