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::hash::Hash;
21use std::mem::size_of_val;
22use std::sync::Arc;
23
24use crate::array_agg::ArrayAgg;
25
26use arrow::array::{ArrayRef, AsArray, BooleanArray, LargeStringArray};
27use arrow::datatypes::{DataType, Field, FieldRef};
28use datafusion_common::cast::{as_generic_string_array, as_string_view_array};
29use datafusion_common::{
30    Result, ScalarValue, internal_datafusion_err, internal_err, not_impl_err,
31};
32use datafusion_expr::function::AccumulatorArgs;
33use datafusion_expr::utils::format_state_name;
34use datafusion_expr::{
35    Accumulator, AggregateUDFImpl, Documentation, EmitTo, GroupsAccumulator, Signature,
36    TypeSignature, Volatility,
37};
38use datafusion_functions_aggregate_common::accumulator::StateFieldsArgs;
39use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::apply_filter_as_nulls;
40use datafusion_macros::user_doc;
41use datafusion_physical_expr::expressions::Literal;
42
43make_udaf_expr_and_func!(
44    StringAgg,
45    string_agg,
46    expr delimiter,
47    "Concatenates the values of string expressions and places separator values between them",
48    string_agg_udaf
49);
50
51#[user_doc(
52    doc_section(label = "General Functions"),
53    description = "Concatenates the values of string expressions and places separator values between them. \
54If ordering is required, strings are concatenated in the specified order. \
55This aggregation function can only mix DISTINCT and ORDER BY if the ordering expression is exactly the same as the first argument expression.",
56    syntax_example = "string_agg([DISTINCT] expression, delimiter [ORDER BY expression])",
57    sql_example = r#"```sql
58> SELECT string_agg(name, ', ') AS names_list
59  FROM employee;
60+--------------------------+
61| names_list               |
62+--------------------------+
63| Alice, Bob, Bob, Charlie |
64+--------------------------+
65> SELECT string_agg(name, ', ' ORDER BY name DESC) AS names_list
66  FROM employee;
67+--------------------------+
68| names_list               |
69+--------------------------+
70| Charlie, Bob, Bob, Alice |
71+--------------------------+
72> SELECT string_agg(DISTINCT name, ', ' ORDER BY name DESC) AS names_list
73  FROM employee;
74+--------------------------+
75| names_list               |
76+--------------------------+
77| Charlie, Bob, Alice |
78+--------------------------+
79```"#,
80    argument(
81        name = "expression",
82        description = "The string expression to concatenate. Can be a column or any valid string expression."
83    ),
84    argument(
85        name = "delimiter",
86        description = "A literal string used as a separator between the concatenated values."
87    )
88)]
89/// STRING_AGG aggregate expression
90#[derive(Debug, PartialEq, Eq, Hash)]
91pub struct StringAgg {
92    signature: Signature,
93    array_agg: ArrayAgg,
94}
95
96impl StringAgg {
97    /// Create a new StringAgg aggregate function
98    pub fn new() -> Self {
99        Self {
100            signature: Signature::one_of(
101                vec![
102                    TypeSignature::Exact(vec![DataType::LargeUtf8, DataType::Utf8]),
103                    TypeSignature::Exact(vec![DataType::LargeUtf8, DataType::LargeUtf8]),
104                    TypeSignature::Exact(vec![DataType::LargeUtf8, DataType::Null]),
105                    TypeSignature::Exact(vec![DataType::LargeUtf8, DataType::Utf8View]),
106                    TypeSignature::Exact(vec![DataType::Utf8, DataType::Utf8]),
107                    TypeSignature::Exact(vec![DataType::Utf8, DataType::LargeUtf8]),
108                    TypeSignature::Exact(vec![DataType::Utf8, DataType::Null]),
109                    TypeSignature::Exact(vec![DataType::Utf8, DataType::Utf8View]),
110                    TypeSignature::Exact(vec![DataType::Utf8View, DataType::Utf8View]),
111                    TypeSignature::Exact(vec![DataType::Utf8View, DataType::LargeUtf8]),
112                    TypeSignature::Exact(vec![DataType::Utf8View, DataType::Null]),
113                    TypeSignature::Exact(vec![DataType::Utf8View, DataType::Utf8]),
114                ],
115                Volatility::Immutable,
116            ),
117            array_agg: Default::default(),
118        }
119    }
120
121    /// Extract the delimiter string from the second argument expression.
122    fn extract_delimiter(args: &AccumulatorArgs) -> Result<String> {
123        let Some(lit) = args.exprs[1].downcast_ref::<Literal>() else {
124            return not_impl_err!("string_agg delimiter must be a string literal");
125        };
126
127        if lit.value().is_null() {
128            return Ok(String::new());
129        }
130
131        match lit.value().try_as_str() {
132            Some(s) => Ok(s.unwrap_or("").to_string()),
133            None => {
134                not_impl_err!(
135                    "string_agg not supported for delimiter \"{}\"",
136                    lit.value()
137                )
138            }
139        }
140    }
141}
142
143impl Default for StringAgg {
144    fn default() -> Self {
145        Self::new()
146    }
147}
148
149/// Three accumulation strategies depending on query shape:
150/// - No DISTINCT / ORDER BY with GROUP BY: `StringAggGroupsAccumulator`
151/// - No DISTINCT / ORDER BY without GROUP BY: `SimpleStringAggAccumulator`
152/// - With DISTINCT or ORDER BY: `StringAggAccumulator` (delegates to `ArrayAgg`)
153impl AggregateUDFImpl for StringAgg {
154    fn name(&self) -> &str {
155        "string_agg"
156    }
157
158    fn signature(&self) -> &Signature {
159        &self.signature
160    }
161
162    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
163        Ok(DataType::LargeUtf8)
164    }
165
166    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
167        if !args.is_distinct && args.ordering_fields.is_empty() {
168            Ok(vec![
169                Field::new(
170                    format_state_name(args.name, "string_agg"),
171                    DataType::LargeUtf8,
172                    true,
173                )
174                .into(),
175            ])
176        } else {
177            self.array_agg.state_fields(args)
178        }
179    }
180
181    fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
182        let delimiter = Self::extract_delimiter(&acc_args)?;
183
184        if !acc_args.is_distinct && acc_args.order_bys.is_empty() {
185            Ok(Box::new(SimpleStringAggAccumulator::new(&delimiter)))
186        } else {
187            let array_agg_acc = self.array_agg.accumulator(AccumulatorArgs {
188                return_field: Field::new(
189                    "f",
190                    DataType::new_list(acc_args.return_field.data_type().clone(), true),
191                    true,
192                )
193                .into(),
194                exprs: &filter_index(acc_args.exprs, 1),
195                expr_fields: &filter_index(acc_args.expr_fields, 1),
196                // Unchanged below; we list each field explicitly in case we ever add more
197                // fields to AccumulatorArgs making it easier to see if changes are also
198                // needed here.
199                schema: acc_args.schema,
200                ignore_nulls: acc_args.ignore_nulls,
201                order_bys: acc_args.order_bys,
202                is_reversed: acc_args.is_reversed,
203                name: acc_args.name,
204                is_distinct: acc_args.is_distinct,
205            })?;
206
207            Ok(Box::new(StringAggAccumulator::new(
208                array_agg_acc,
209                &delimiter,
210            )))
211        }
212    }
213
214    fn reverse_expr(&self) -> datafusion_expr::ReversedUDAF {
215        datafusion_expr::ReversedUDAF::Reversed(string_agg_udaf())
216    }
217
218    fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
219        !args.is_distinct && args.order_bys.is_empty()
220    }
221
222    fn create_groups_accumulator(
223        &self,
224        args: AccumulatorArgs,
225    ) -> Result<Box<dyn GroupsAccumulator>> {
226        let delimiter = Self::extract_delimiter(&args)?;
227        Ok(Box::new(StringAggGroupsAccumulator::new(delimiter)))
228    }
229
230    fn documentation(&self) -> Option<&Documentation> {
231        self.doc()
232    }
233}
234
235/// StringAgg accumulator for the general case (with order or distinct specified)
236#[derive(Debug)]
237pub(crate) struct StringAggAccumulator {
238    array_agg_acc: Box<dyn Accumulator>,
239    delimiter: String,
240}
241
242impl StringAggAccumulator {
243    pub fn new(array_agg_acc: Box<dyn Accumulator>, delimiter: &str) -> Self {
244        Self {
245            array_agg_acc,
246            delimiter: delimiter.to_string(),
247        }
248    }
249}
250
251impl Accumulator for StringAggAccumulator {
252    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
253        self.array_agg_acc.update_batch(&filter_index(values, 1))
254    }
255
256    fn evaluate(&mut self) -> Result<ScalarValue> {
257        let scalar = self.array_agg_acc.evaluate()?;
258
259        let ScalarValue::List(list) = scalar else {
260            return internal_err!(
261                "Expected a DataType::List while evaluating underlying ArrayAggAccumulator, but got {}",
262                scalar.data_type()
263            );
264        };
265
266        let string_arr: Vec<_> = match list.value_type() {
267            DataType::LargeUtf8 => as_generic_string_array::<i64>(list.values())?
268                .iter()
269                .flatten()
270                .collect(),
271            DataType::Utf8 => as_generic_string_array::<i32>(list.values())?
272                .iter()
273                .flatten()
274                .collect(),
275            DataType::Utf8View => as_string_view_array(list.values())?
276                .iter()
277                .flatten()
278                .collect(),
279            _ => {
280                return internal_err!(
281                    "Expected elements to of type Utf8 or LargeUtf8, but got {}",
282                    list.value_type()
283                );
284            }
285        };
286
287        if string_arr.is_empty() {
288            return Ok(ScalarValue::LargeUtf8(None));
289        }
290
291        Ok(ScalarValue::LargeUtf8(Some(
292            string_arr.join(&self.delimiter),
293        )))
294    }
295
296    fn size(&self) -> usize {
297        size_of_val(self) - size_of_val(&self.array_agg_acc)
298            + self.array_agg_acc.size()
299            + self.delimiter.capacity()
300    }
301
302    fn state(&mut self) -> Result<Vec<ScalarValue>> {
303        self.array_agg_acc.state()
304    }
305
306    fn merge_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
307        self.array_agg_acc.merge_batch(values)
308    }
309}
310
311fn filter_index<T: Clone>(values: &[T], index: usize) -> Vec<T> {
312    values
313        .iter()
314        .enumerate()
315        .filter(|(i, _)| *i != index)
316        .map(|(_, v)| v)
317        .cloned()
318        .collect::<Vec<_>>()
319}
320
321/// GroupsAccumulator for `string_agg` without DISTINCT or ORDER BY.
322#[derive(Debug)]
323struct StringAggGroupsAccumulator {
324    /// The delimiter placed between concatenated values.
325    delimiter: String,
326    /// Accumulated string per group. `None` means no values have been seen
327    /// (the group's output will be NULL).
328    /// A potential improvement is to avoid this String allocation
329    /// See <https://github.com/apache/datafusion/issues/21156>
330    values: Vec<Option<String>>,
331    /// Running total of string data bytes across all groups.
332    total_data_bytes: usize,
333}
334
335impl StringAggGroupsAccumulator {
336    fn new(delimiter: String) -> Self {
337        Self {
338            delimiter,
339            values: Vec::new(),
340            total_data_bytes: 0,
341        }
342    }
343
344    fn append_batch<'a>(
345        &mut self,
346        iter: impl Iterator<Item = Option<&'a str>>,
347        group_indices: &[usize],
348    ) {
349        for (opt_value, &group_idx) in iter.zip(group_indices.iter()) {
350            if let Some(value) = opt_value {
351                match &mut self.values[group_idx] {
352                    Some(existing) => {
353                        let added = self.delimiter.len() + value.len();
354                        existing.reserve(added);
355                        existing.push_str(&self.delimiter);
356                        existing.push_str(value);
357                        self.total_data_bytes += added;
358                    }
359                    slot @ None => {
360                        *slot = Some(value.to_string());
361                        self.total_data_bytes += value.len();
362                    }
363                }
364            }
365        }
366    }
367}
368
369impl GroupsAccumulator for StringAggGroupsAccumulator {
370    fn update_batch(
371        &mut self,
372        values: &[ArrayRef],
373        group_indices: &[usize],
374        opt_filter: Option<&BooleanArray>,
375        total_num_groups: usize,
376    ) -> Result<()> {
377        self.values.resize(total_num_groups, None);
378        let array = apply_filter_as_nulls(&values[0], opt_filter)?;
379        match array.data_type() {
380            DataType::Utf8 => {
381                self.append_batch(array.as_string::<i32>().iter(), group_indices)
382            }
383            DataType::LargeUtf8 => {
384                self.append_batch(array.as_string::<i64>().iter(), group_indices)
385            }
386            DataType::Utf8View => {
387                self.append_batch(array.as_string_view().iter(), group_indices)
388            }
389            other => {
390                return internal_err!("string_agg unexpected data type: {other}");
391            }
392        }
393        Ok(())
394    }
395
396    fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
397        let to_emit = emit_to.take_needed(&mut self.values);
398        let emitted_bytes: usize = to_emit
399            .iter()
400            .filter_map(|opt| opt.as_ref().map(|s| s.len()))
401            .sum();
402        self.total_data_bytes -= emitted_bytes;
403
404        let result: ArrayRef = Arc::new(LargeStringArray::from(to_emit));
405        Ok(result)
406    }
407
408    fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
409        self.evaluate(emit_to).map(|arr| vec![arr])
410    }
411
412    fn merge_batch(
413        &mut self,
414        values: &[ArrayRef],
415        group_indices: &[usize],
416        total_num_groups: usize,
417    ) -> Result<()> {
418        // State is always LargeUtf8, which update_batch already handles.
419        self.update_batch(values, group_indices, None, total_num_groups)
420    }
421
422    fn convert_to_state(
423        &self,
424        values: &[ArrayRef],
425        opt_filter: Option<&BooleanArray>,
426    ) -> Result<Vec<ArrayRef>> {
427        let input = apply_filter_as_nulls(&values[0], opt_filter)?;
428        let result = if input.data_type() == &DataType::LargeUtf8 {
429            input
430        } else {
431            arrow::compute::cast(&input, &DataType::LargeUtf8)?
432        };
433        Ok(vec![result])
434    }
435    fn size(&self) -> usize {
436        self.total_data_bytes
437            + self.values.capacity() * size_of::<Option<String>>()
438            + self.delimiter.capacity()
439            + size_of_val(self)
440    }
441}
442
443/// Per-row accumulator for `string_agg` without DISTINCT or ORDER BY.  Used for
444/// non-grouped aggregation; grouped queries use [`StringAggGroupsAccumulator`].
445#[derive(Debug)]
446pub(crate) struct SimpleStringAggAccumulator {
447    delimiter: String,
448    /// Updated during `update_batch()`. e.g. "foo,bar"
449    accumulated_string: String,
450    has_value: bool,
451}
452
453impl SimpleStringAggAccumulator {
454    pub fn new(delimiter: &str) -> Self {
455        Self {
456            delimiter: delimiter.to_string(),
457            accumulated_string: String::new(),
458            has_value: false,
459        }
460    }
461
462    #[inline]
463    fn append_strings<'a, I>(&mut self, iter: I)
464    where
465        I: Iterator<Item = Option<&'a str>>,
466    {
467        for value in iter.flatten() {
468            if self.has_value {
469                self.accumulated_string.push_str(&self.delimiter);
470            }
471
472            self.accumulated_string.push_str(value);
473            self.has_value = true;
474        }
475    }
476}
477
478impl Accumulator for SimpleStringAggAccumulator {
479    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
480        let string_arr = values.first().ok_or_else(|| {
481            internal_datafusion_err!(
482                "Planner should ensure its first arg is Utf8/Utf8View"
483            )
484        })?;
485
486        match string_arr.data_type() {
487            DataType::Utf8 => self.append_strings(string_arr.as_string::<i32>().iter()),
488            DataType::LargeUtf8 => {
489                self.append_strings(string_arr.as_string::<i64>().iter())
490            }
491            DataType::Utf8View => self.append_strings(string_arr.as_string_view().iter()),
492            other => {
493                return internal_err!(
494                    "Planner should ensure string_agg first argument is Utf8-like, found {other}"
495                );
496            }
497        }
498
499        Ok(())
500    }
501
502    fn evaluate(&mut self) -> Result<ScalarValue> {
503        if self.has_value {
504            Ok(ScalarValue::LargeUtf8(Some(
505                self.accumulated_string.clone(),
506            )))
507        } else {
508            Ok(ScalarValue::LargeUtf8(None))
509        }
510    }
511
512    fn size(&self) -> usize {
513        size_of_val(self) + self.delimiter.capacity() + self.accumulated_string.capacity()
514    }
515
516    fn state(&mut self) -> Result<Vec<ScalarValue>> {
517        let result = if self.has_value {
518            ScalarValue::LargeUtf8(Some(std::mem::take(&mut self.accumulated_string)))
519        } else {
520            ScalarValue::LargeUtf8(None)
521        };
522        self.has_value = false;
523
524        Ok(vec![result])
525    }
526
527    fn merge_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
528        self.update_batch(values)
529    }
530}
531
532#[cfg(test)]
533mod tests {
534    use super::*;
535    use arrow::array::LargeStringArray;
536    use arrow::compute::SortOptions;
537    use arrow::datatypes::{Fields, Schema};
538    use datafusion_physical_expr::expressions::Column;
539    use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
540    use std::sync::Arc;
541
542    #[test]
543    fn no_duplicates_no_distinct() -> Result<()> {
544        let (mut acc1, mut acc2) = StringAggAccumulatorBuilder::new(",").build_two()?;
545
546        acc1.update_batch(&[data(["a", "b", "c"]), data([","])])?;
547        acc2.update_batch(&[data(["d", "e", "f"]), data([","])])?;
548        acc1 = merge(acc1, acc2)?;
549
550        let result = some_str(acc1.evaluate()?);
551
552        assert_eq!(result, "a,b,c,d,e,f");
553
554        Ok(())
555    }
556
557    #[test]
558    fn no_duplicates_distinct() -> Result<()> {
559        let (mut acc1, mut acc2) = StringAggAccumulatorBuilder::new(",")
560            .distinct()
561            .build_two()?;
562
563        acc1.update_batch(&[data(["a", "b", "c"]), data([","])])?;
564        acc2.update_batch(&[data(["d", "e", "f"]), data([","])])?;
565        acc1 = merge(acc1, acc2)?;
566
567        let result = some_str_sorted(acc1.evaluate()?, ",");
568
569        assert_eq!(result, "a,b,c,d,e,f");
570
571        Ok(())
572    }
573
574    #[test]
575    fn duplicates_no_distinct() -> Result<()> {
576        let (mut acc1, mut acc2) = StringAggAccumulatorBuilder::new(",").build_two()?;
577
578        acc1.update_batch(&[data(["a", "b", "c"]), data([","])])?;
579        acc2.update_batch(&[data(["a", "b", "c"]), data([","])])?;
580        acc1 = merge(acc1, acc2)?;
581
582        let result = some_str(acc1.evaluate()?);
583
584        assert_eq!(result, "a,b,c,a,b,c");
585
586        Ok(())
587    }
588
589    #[test]
590    fn duplicates_distinct() -> Result<()> {
591        let (mut acc1, mut acc2) = StringAggAccumulatorBuilder::new(",")
592            .distinct()
593            .build_two()?;
594
595        acc1.update_batch(&[data(["a", "b", "c"]), data([","])])?;
596        acc2.update_batch(&[data(["a", "b", "c"]), data([","])])?;
597        acc1 = merge(acc1, acc2)?;
598
599        let result = some_str_sorted(acc1.evaluate()?, ",");
600
601        assert_eq!(result, "a,b,c");
602
603        Ok(())
604    }
605
606    #[test]
607    fn no_duplicates_distinct_sort_asc() -> Result<()> {
608        let (mut acc1, mut acc2) = StringAggAccumulatorBuilder::new(",")
609            .distinct()
610            .order_by_col("col", SortOptions::new(false, false))
611            .build_two()?;
612
613        acc1.update_batch(&[data(["e", "b", "d"]), data([","])])?;
614        acc2.update_batch(&[data(["f", "a", "c"]), data([","])])?;
615        acc1 = merge(acc1, acc2)?;
616
617        let result = some_str(acc1.evaluate()?);
618
619        assert_eq!(result, "a,b,c,d,e,f");
620
621        Ok(())
622    }
623
624    #[test]
625    fn no_duplicates_distinct_sort_desc() -> Result<()> {
626        let (mut acc1, mut acc2) = StringAggAccumulatorBuilder::new(",")
627            .distinct()
628            .order_by_col("col", SortOptions::new(true, false))
629            .build_two()?;
630
631        acc1.update_batch(&[data(["e", "b", "d"]), data([","])])?;
632        acc2.update_batch(&[data(["f", "a", "c"]), data([","])])?;
633        acc1 = merge(acc1, acc2)?;
634
635        let result = some_str(acc1.evaluate()?);
636
637        assert_eq!(result, "f,e,d,c,b,a");
638
639        Ok(())
640    }
641
642    #[test]
643    fn duplicates_distinct_sort_asc() -> Result<()> {
644        let (mut acc1, mut acc2) = StringAggAccumulatorBuilder::new(",")
645            .distinct()
646            .order_by_col("col", SortOptions::new(false, false))
647            .build_two()?;
648
649        acc1.update_batch(&[data(["a", "c", "b"]), data([","])])?;
650        acc2.update_batch(&[data(["b", "c", "a"]), data([","])])?;
651        acc1 = merge(acc1, acc2)?;
652
653        let result = some_str(acc1.evaluate()?);
654
655        assert_eq!(result, "a,b,c");
656
657        Ok(())
658    }
659
660    #[test]
661    fn duplicates_distinct_sort_desc() -> Result<()> {
662        let (mut acc1, mut acc2) = StringAggAccumulatorBuilder::new(",")
663            .distinct()
664            .order_by_col("col", SortOptions::new(true, false))
665            .build_two()?;
666
667        acc1.update_batch(&[data(["a", "c", "b"]), data([","])])?;
668        acc2.update_batch(&[data(["b", "c", "a"]), data([","])])?;
669        acc1 = merge(acc1, acc2)?;
670
671        let result = some_str(acc1.evaluate()?);
672
673        assert_eq!(result, "c,b,a");
674
675        Ok(())
676    }
677
678    struct StringAggAccumulatorBuilder {
679        sep: String,
680        distinct: bool,
681        order_bys: Vec<PhysicalSortExpr>,
682        schema: Schema,
683    }
684
685    impl StringAggAccumulatorBuilder {
686        fn new(sep: &str) -> Self {
687            Self {
688                sep: sep.to_string(),
689                distinct: Default::default(),
690                order_bys: vec![],
691                schema: Schema {
692                    fields: Fields::from(vec![Field::new(
693                        "col",
694                        DataType::LargeUtf8,
695                        true,
696                    )]),
697                    metadata: Default::default(),
698                },
699            }
700        }
701        fn distinct(mut self) -> Self {
702            self.distinct = true;
703            self
704        }
705
706        fn order_by_col(mut self, col: &str, sort_options: SortOptions) -> Self {
707            self.order_bys.extend([PhysicalSortExpr::new(
708                Arc::new(
709                    Column::new_with_schema(col, &self.schema)
710                        .expect("column not available in schema"),
711                ),
712                sort_options,
713            )]);
714            self
715        }
716
717        fn build(&self) -> Result<Box<dyn Accumulator>> {
718            StringAgg::new().accumulator(AccumulatorArgs {
719                return_field: Field::new("f", DataType::LargeUtf8, true).into(),
720                schema: &self.schema,
721                expr_fields: &[
722                    Field::new("col", DataType::LargeUtf8, true).into(),
723                    Field::new("lit", DataType::Utf8, false).into(),
724                ],
725                ignore_nulls: false,
726                order_bys: &self.order_bys,
727                is_reversed: false,
728                name: "",
729                is_distinct: self.distinct,
730                exprs: &[
731                    Arc::new(Column::new("col", 0)),
732                    Arc::new(Literal::new(ScalarValue::Utf8(Some(self.sep.to_string())))),
733                ],
734            })
735        }
736
737        fn build_two(&self) -> Result<(Box<dyn Accumulator>, Box<dyn Accumulator>)> {
738            Ok((self.build()?, self.build()?))
739        }
740    }
741
742    fn some_str(value: ScalarValue) -> String {
743        str(value)
744            .expect("ScalarValue was not a String")
745            .expect("ScalarValue was None")
746    }
747
748    fn some_str_sorted(value: ScalarValue, sep: &str) -> String {
749        let value = some_str(value);
750        let mut parts: Vec<&str> = value.split(sep).collect();
751        parts.sort();
752        parts.join(sep)
753    }
754
755    fn str(value: ScalarValue) -> Result<Option<String>> {
756        match value {
757            ScalarValue::LargeUtf8(v) => Ok(v),
758            _ => internal_err!(
759                "Expected ScalarValue::LargeUtf8, got {}",
760                value.data_type()
761            ),
762        }
763    }
764
765    fn data<const N: usize>(list: [&str; N]) -> ArrayRef {
766        Arc::new(LargeStringArray::from(list.to_vec()))
767    }
768
769    fn merge(
770        mut acc1: Box<dyn Accumulator>,
771        mut acc2: Box<dyn Accumulator>,
772    ) -> Result<Box<dyn Accumulator>> {
773        let intermediate_state = acc2.state().and_then(|e| {
774            e.iter()
775                .map(|v| v.to_array())
776                .collect::<Result<Vec<ArrayRef>>>()
777        })?;
778        acc1.merge_batch(&intermediate_state)?;
779        Ok(acc1)
780    }
781
782    // ---------------------------------------------------------------
783    // Tests for StringAggGroupsAccumulator
784    // ---------------------------------------------------------------
785
786    fn make_groups_acc(delimiter: &str) -> StringAggGroupsAccumulator {
787        StringAggGroupsAccumulator::new(delimiter.to_string())
788    }
789
790    /// Helper: evaluate and downcast to LargeStringArray
791    fn evaluate_groups(
792        acc: &mut StringAggGroupsAccumulator,
793        emit_to: EmitTo,
794    ) -> Vec<Option<String>> {
795        let result = acc.evaluate(emit_to).unwrap();
796        let arr = result.as_any().downcast_ref::<LargeStringArray>().unwrap();
797        arr.iter().map(|v| v.map(|s| s.to_string())).collect()
798    }
799
800    #[test]
801    fn groups_basic() -> Result<()> {
802        let mut acc = make_groups_acc(",");
803
804        // 6 rows, 3 groups: group 0 gets "a","d"; group 1 gets "b","e"; group 2 gets "c","f"
805        let values: ArrayRef =
806            Arc::new(LargeStringArray::from(vec!["a", "b", "c", "d", "e", "f"]));
807        let group_indices = vec![0, 1, 2, 0, 1, 2];
808        acc.update_batch(&[values], &group_indices, None, 3)?;
809
810        let result = evaluate_groups(&mut acc, EmitTo::All);
811        assert_eq!(
812            result,
813            vec![
814                Some("a,d".to_string()),
815                Some("b,e".to_string()),
816                Some("c,f".to_string()),
817            ]
818        );
819        Ok(())
820    }
821
822    #[test]
823    fn groups_with_nulls() -> Result<()> {
824        let mut acc = make_groups_acc("|");
825
826        // Group 0: "a", NULL, "c" → "a|c"
827        // Group 1: NULL, "b"     → "b"
828        // Group 2: NULL only     → NULL
829        let values: ArrayRef = Arc::new(LargeStringArray::from(vec![
830            Some("a"),
831            None,
832            Some("c"),
833            None,
834            Some("b"),
835            None,
836        ]));
837        let group_indices = vec![0, 1, 0, 2, 1, 2];
838        acc.update_batch(&[values], &group_indices, None, 3)?;
839
840        let result = evaluate_groups(&mut acc, EmitTo::All);
841        assert_eq!(
842            result,
843            vec![Some("a|c".to_string()), Some("b".to_string()), None,]
844        );
845        Ok(())
846    }
847
848    #[test]
849    fn groups_with_filter() -> Result<()> {
850        let mut acc = make_groups_acc(",");
851
852        let values: ArrayRef = Arc::new(LargeStringArray::from(vec!["a", "b", "c", "d"]));
853        let group_indices = vec![0, 0, 1, 1];
854        // Filter: only rows 0 and 3 are included
855        let filter = BooleanArray::from(vec![true, false, false, true]);
856        acc.update_batch(&[values], &group_indices, Some(&filter), 2)?;
857
858        let result = evaluate_groups(&mut acc, EmitTo::All);
859        assert_eq!(result, vec![Some("a".to_string()), Some("d".to_string())]);
860        Ok(())
861    }
862
863    #[test]
864    fn groups_emit_first() -> Result<()> {
865        let mut acc = make_groups_acc(",");
866
867        let values: ArrayRef =
868            Arc::new(LargeStringArray::from(vec!["a", "b", "c", "d", "e", "f"]));
869        let group_indices = vec![0, 1, 2, 0, 1, 2];
870        acc.update_batch(&[values], &group_indices, None, 3)?;
871
872        // Emit only the first 2 groups
873        let result = evaluate_groups(&mut acc, EmitTo::First(2));
874        assert_eq!(
875            result,
876            vec![Some("a,d".to_string()), Some("b,e".to_string())]
877        );
878
879        // Group 2 (now shifted to index 0) should still be intact
880        let result = evaluate_groups(&mut acc, EmitTo::All);
881        assert_eq!(result, vec![Some("c,f".to_string())]);
882        Ok(())
883    }
884
885    #[test]
886    fn groups_merge_batch() -> Result<()> {
887        let mut acc = make_groups_acc(",");
888
889        // First batch: group 0 = "a", group 1 = "b"
890        let values: ArrayRef = Arc::new(LargeStringArray::from(vec!["a", "b"]));
891        acc.update_batch(&[values], &[0, 1], None, 2)?;
892
893        // Simulate a second accumulator's state (LargeUtf8 partial strings)
894        let partial_state: ArrayRef = Arc::new(LargeStringArray::from(vec!["c,d", "e"]));
895        acc.merge_batch(&[partial_state], &[0, 1], 2)?;
896
897        let result = evaluate_groups(&mut acc, EmitTo::All);
898        assert_eq!(
899            result,
900            vec![Some("a,c,d".to_string()), Some("b,e".to_string())]
901        );
902        Ok(())
903    }
904
905    #[test]
906    fn groups_empty_groups() -> Result<()> {
907        let mut acc = make_groups_acc(",");
908
909        // 4 groups total, but only groups 0 and 2 receive values
910        let values: ArrayRef = Arc::new(LargeStringArray::from(vec!["a", "b"]));
911        acc.update_batch(&[values], &[0, 2], None, 4)?;
912
913        let result = evaluate_groups(&mut acc, EmitTo::All);
914        assert_eq!(
915            result,
916            vec![
917                Some("a".to_string()),
918                None, // group 1: never received a value
919                Some("b".to_string()),
920                None, // group 3: never received a value
921            ]
922        );
923        Ok(())
924    }
925
926    #[test]
927    fn groups_multiple_batches() -> Result<()> {
928        let mut acc = make_groups_acc("|");
929
930        // Batch 1: 2 groups
931        let values: ArrayRef = Arc::new(LargeStringArray::from(vec!["a", "b"]));
932        acc.update_batch(&[values], &[0, 1], None, 2)?;
933
934        // Batch 2: same groups, plus a new group
935        let values: ArrayRef = Arc::new(LargeStringArray::from(vec!["c", "d", "e"]));
936        acc.update_batch(&[values], &[0, 1, 2], None, 3)?;
937
938        let result = evaluate_groups(&mut acc, EmitTo::All);
939        assert_eq!(
940            result,
941            vec![
942                Some("a|c".to_string()),
943                Some("b|d".to_string()),
944                Some("e".to_string()),
945            ]
946        );
947        Ok(())
948    }
949}
datafusion_functions_aggregate/string_agg.rs

datafusion_functions_aggregate/
string_agg.rs
