1use 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#[derive(Debug, PartialEq, Eq, Hash)]
86pub struct StringAgg {
87 signature: Signature,
88 array_agg: ArrayAgg,
89}
90
91impl StringAgg {
92 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 ""
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}