datafusion_functions/string/
concat_ws.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::{as_largestring_array, Array, StringArray};
19use std::any::Any;
20use std::sync::Arc;
21
22use arrow::datatypes::DataType;
23
24use crate::string::concat;
25use crate::string::concat::simplify_concat;
26use crate::string::concat_ws;
27use crate::strings::{ColumnarValueRef, StringArrayBuilder};
28use datafusion_common::cast::{as_string_array, as_string_view_array};
29use datafusion_common::{exec_err, internal_err, plan_err, Result, ScalarValue};
30use datafusion_expr::expr::ScalarFunction;
31use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
32use datafusion_expr::{lit, ColumnarValue, Documentation, Expr, Volatility};
33use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature};
34use datafusion_macros::user_doc;
35
36#[user_doc(
37    doc_section(label = "String Functions"),
38    description = "Concatenates multiple strings together with a specified separator.",
39    syntax_example = "concat_ws(separator, str[, ..., str_n])",
40    sql_example = r#"```sql
41> select concat_ws('_', 'data', 'fusion');
42+--------------------------------------------------+
43| concat_ws(Utf8("_"),Utf8("data"),Utf8("fusion")) |
44+--------------------------------------------------+
45| data_fusion                                      |
46+--------------------------------------------------+
47```"#,
48    argument(
49        name = "separator",
50        description = "Separator to insert between concatenated strings."
51    ),
52    argument(
53        name = "str",
54        description = "String expression to operate on. Can be a constant, column, or function, and any combination of operators."
55    ),
56    argument(
57        name = "str_n",
58        description = "Subsequent string expressions to concatenate."
59    ),
60    related_udf(name = "concat")
61)]
62#[derive(Debug)]
63pub struct ConcatWsFunc {
64    signature: Signature,
65}
66
67impl Default for ConcatWsFunc {
68    fn default() -> Self {
69        ConcatWsFunc::new()
70    }
71}
72
73impl ConcatWsFunc {
74    pub fn new() -> Self {
75        use DataType::*;
76        Self {
77            signature: Signature::variadic(
78                vec![Utf8View, Utf8, LargeUtf8],
79                Volatility::Immutable,
80            ),
81        }
82    }
83}
84
85impl ScalarUDFImpl for ConcatWsFunc {
86    fn as_any(&self) -> &dyn Any {
87        self
88    }
89
90    fn name(&self) -> &str {
91        "concat_ws"
92    }
93
94    fn signature(&self) -> &Signature {
95        &self.signature
96    }
97
98    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
99        use DataType::*;
100        Ok(Utf8)
101    }
102
103    /// Concatenates all but the first argument, with separators. The first argument is used as the separator string, and should not be NULL. Other NULL arguments are ignored.
104    /// concat_ws(',', 'abcde', 2, NULL, 22) = 'abcde,2,22'
105    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
106        let ScalarFunctionArgs { args, .. } = args;
107
108        // do not accept 0 arguments.
109        if args.len() < 2 {
110            return exec_err!(
111                "concat_ws was called with {} arguments. It requires at least 2.",
112                args.len()
113            );
114        }
115
116        let array_len = args
117            .iter()
118            .filter_map(|x| match x {
119                ColumnarValue::Array(array) => Some(array.len()),
120                _ => None,
121            })
122            .next();
123
124        // Scalar
125        if array_len.is_none() {
126            let ColumnarValue::Scalar(scalar) = &args[0] else {
127                // loop above checks for all args being scalar
128                unreachable!()
129            };
130            let sep = match scalar.try_as_str() {
131                Some(Some(s)) => s,
132                Some(None) => {
133                    // null literal string
134                    return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None)));
135                }
136                None => return internal_err!("Expected string literal, got {scalar:?}"),
137            };
138
139            let mut result = String::new();
140            // iterator over Option<str>
141            let iter = &mut args[1..].iter().map(|arg| {
142                let ColumnarValue::Scalar(scalar) = arg else {
143                    // loop above checks for all args being scalar
144                    unreachable!()
145                };
146                scalar.try_as_str()
147            });
148
149            // append first non null arg
150            for scalar in iter.by_ref() {
151                match scalar {
152                    Some(Some(s)) => {
153                        result.push_str(s);
154                        break;
155                    }
156                    Some(None) => {} // null literal string
157                    None => {
158                        return internal_err!("Expected string literal, got {scalar:?}")
159                    }
160                }
161            }
162
163            // handle subsequent non null args
164            for scalar in iter.by_ref() {
165                match scalar {
166                    Some(Some(s)) => {
167                        result.push_str(sep);
168                        result.push_str(s);
169                    }
170                    Some(None) => {} // null literal string
171                    None => {
172                        return internal_err!("Expected string literal, got {scalar:?}")
173                    }
174                }
175            }
176
177            return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(result))));
178        }
179
180        // Array
181        let len = array_len.unwrap();
182        let mut data_size = 0;
183
184        // parse sep
185        let sep = match &args[0] {
186            ColumnarValue::Scalar(ScalarValue::Utf8(Some(s))) => {
187                data_size += s.len() * len * (args.len() - 2); // estimate
188                ColumnarValueRef::Scalar(s.as_bytes())
189            }
190            ColumnarValue::Scalar(ScalarValue::Utf8(None)) => {
191                return Ok(ColumnarValue::Array(Arc::new(StringArray::new_null(len))));
192            }
193            ColumnarValue::Array(array) => {
194                let string_array = as_string_array(array)?;
195                data_size += string_array.values().len() * (args.len() - 2); // estimate
196                if array.is_nullable() {
197                    ColumnarValueRef::NullableArray(string_array)
198                } else {
199                    ColumnarValueRef::NonNullableArray(string_array)
200                }
201            }
202            _ => unreachable!("concat ws"),
203        };
204
205        let mut columns = Vec::with_capacity(args.len() - 1);
206        for arg in &args[1..] {
207            match arg {
208                ColumnarValue::Scalar(ScalarValue::Utf8(maybe_value))
209                | ColumnarValue::Scalar(ScalarValue::LargeUtf8(maybe_value))
210                | ColumnarValue::Scalar(ScalarValue::Utf8View(maybe_value)) => {
211                    if let Some(s) = maybe_value {
212                        data_size += s.len() * len;
213                        columns.push(ColumnarValueRef::Scalar(s.as_bytes()));
214                    }
215                }
216                ColumnarValue::Array(array) => {
217                    match array.data_type() {
218                        DataType::Utf8 => {
219                            let string_array = as_string_array(array)?;
220
221                            data_size += string_array.values().len();
222                            let column = if array.is_nullable() {
223                                ColumnarValueRef::NullableArray(string_array)
224                            } else {
225                                ColumnarValueRef::NonNullableArray(string_array)
226                            };
227                            columns.push(column);
228                        },
229                        DataType::LargeUtf8 => {
230                            let string_array = as_largestring_array(array);
231
232                            data_size += string_array.values().len();
233                            let column = if array.is_nullable() {
234                                ColumnarValueRef::NullableLargeStringArray(string_array)
235                            } else {
236                                ColumnarValueRef::NonNullableLargeStringArray(string_array)
237                            };
238                            columns.push(column);
239                        },
240                        DataType::Utf8View => {
241                            let string_array = as_string_view_array(array)?;
242
243                            data_size += string_array.data_buffers().iter().map(|buf| buf.len()).sum::<usize>();
244                            let column = if array.is_nullable() {
245                                ColumnarValueRef::NullableStringViewArray(string_array)
246                            } else {
247                                ColumnarValueRef::NonNullableStringViewArray(string_array)
248                            };
249                            columns.push(column);
250                        },
251                        other => {
252                            return plan_err!("Input was {other} which is not a supported datatype for concat_ws function.")
253                        }
254                    };
255                }
256                _ => unreachable!(),
257            }
258        }
259
260        let mut builder = StringArrayBuilder::with_capacity(len, data_size);
261        for i in 0..len {
262            if !sep.is_valid(i) {
263                builder.append_offset();
264                continue;
265            }
266
267            let mut iter = columns.iter();
268            for column in iter.by_ref() {
269                if column.is_valid(i) {
270                    builder.write::<false>(column, i);
271                    break;
272                }
273            }
274
275            for column in iter {
276                if column.is_valid(i) {
277                    builder.write::<false>(&sep, i);
278                    builder.write::<false>(column, i);
279                }
280            }
281
282            builder.append_offset();
283        }
284
285        Ok(ColumnarValue::Array(Arc::new(builder.finish(sep.nulls()))))
286    }
287
288    /// Simply the `concat_ws` function by
289    /// 1. folding to `null` if the delimiter is null
290    /// 2. filtering out `null` arguments
291    /// 3. using `concat` to replace `concat_ws` if the delimiter is an empty string
292    /// 4. concatenating contiguous literals if the delimiter is a literal.
293    fn simplify(
294        &self,
295        args: Vec<Expr>,
296        _info: &dyn SimplifyInfo,
297    ) -> Result<ExprSimplifyResult> {
298        match &args[..] {
299            [delimiter, vals @ ..] => simplify_concat_ws(delimiter, vals),
300            _ => Ok(ExprSimplifyResult::Original(args)),
301        }
302    }
303
304    fn documentation(&self) -> Option<&Documentation> {
305        self.doc()
306    }
307}
308
309fn simplify_concat_ws(delimiter: &Expr, args: &[Expr]) -> Result<ExprSimplifyResult> {
310    match delimiter {
311        Expr::Literal(
312            ScalarValue::Utf8(delimiter)
313            | ScalarValue::LargeUtf8(delimiter)
314            | ScalarValue::Utf8View(delimiter),
315            _,
316        ) => {
317            match delimiter {
318                // when the delimiter is an empty string,
319                // we can use `concat` to replace `concat_ws`
320                Some(delimiter) if delimiter.is_empty() => {
321                    match simplify_concat(args.to_vec())? {
322                        ExprSimplifyResult::Original(_) => {
323                            Ok(ExprSimplifyResult::Simplified(Expr::ScalarFunction(
324                                ScalarFunction {
325                                    func: concat(),
326                                    args: args.to_vec(),
327                                },
328                            )))
329                        }
330                        expr => Ok(expr),
331                    }
332                }
333                Some(delimiter) => {
334                    let mut new_args = Vec::with_capacity(args.len());
335                    new_args.push(lit(delimiter));
336                    let mut contiguous_scalar = None;
337                    for arg in args {
338                        match arg {
339                            // filter out null args
340                            Expr::Literal(ScalarValue::Utf8(None) | ScalarValue::LargeUtf8(None) | ScalarValue::Utf8View(None), _) => {}
341                            Expr::Literal(ScalarValue::Utf8(Some(v)) | ScalarValue::LargeUtf8(Some(v)) | ScalarValue::Utf8View(Some(v)), _) => {
342                                match contiguous_scalar {
343                                    None => contiguous_scalar = Some(v.to_string()),
344                                    Some(mut pre) => {
345                                        pre += delimiter;
346                                        pre += v;
347                                        contiguous_scalar = Some(pre)
348                                    }
349                                }
350                            }
351                            Expr::Literal(s, _) => return internal_err!("The scalar {s} should be casted to string type during the type coercion."),
352                            // If the arg is not a literal, we should first push the current `contiguous_scalar`
353                            // to the `new_args` and reset it to None.
354                            // Then pushing this arg to the `new_args`.
355                            arg => {
356                                if let Some(val) = contiguous_scalar {
357                                    new_args.push(lit(val));
358                                }
359                                new_args.push(arg.clone());
360                                contiguous_scalar = None;
361                            }
362                        }
363                    }
364                    if let Some(val) = contiguous_scalar {
365                        new_args.push(lit(val));
366                    }
367
368                    Ok(ExprSimplifyResult::Simplified(Expr::ScalarFunction(
369                        ScalarFunction {
370                            func: concat_ws(),
371                            args: new_args,
372                        },
373                    )))
374                }
375                // if the delimiter is null, then the value of the whole expression is null.
376                None => Ok(ExprSimplifyResult::Simplified(Expr::Literal(
377                    ScalarValue::Utf8(None),
378                    None,
379                ))),
380            }
381        }
382        Expr::Literal(d, _) => internal_err!(
383            "The scalar {d} should be casted to string type during the type coercion."
384        ),
385        _ => {
386            let mut args = args
387                .iter()
388                .filter(|&x| !is_null(x))
389                .cloned()
390                .collect::<Vec<Expr>>();
391            args.insert(0, delimiter.clone());
392            Ok(ExprSimplifyResult::Original(args))
393        }
394    }
395}
396
397fn is_null(expr: &Expr) -> bool {
398    match expr {
399        Expr::Literal(v, _) => v.is_null(),
400        _ => false,
401    }
402}
403
404#[cfg(test)]
405mod tests {
406    use std::sync::Arc;
407
408    use crate::string::concat_ws::ConcatWsFunc;
409    use arrow::array::{Array, ArrayRef, StringArray};
410    use arrow::datatypes::DataType::Utf8;
411    use arrow::datatypes::Field;
412    use datafusion_common::Result;
413    use datafusion_common::ScalarValue;
414    use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
415
416    use crate::utils::test::test_function;
417
418    #[test]
419    fn test_functions() -> Result<()> {
420        test_function!(
421            ConcatWsFunc::new(),
422            vec![
423                ColumnarValue::Scalar(ScalarValue::from("|")),
424                ColumnarValue::Scalar(ScalarValue::from("aa")),
425                ColumnarValue::Scalar(ScalarValue::from("bb")),
426                ColumnarValue::Scalar(ScalarValue::from("cc")),
427            ],
428            Ok(Some("aa|bb|cc")),
429            &str,
430            Utf8,
431            StringArray
432        );
433        test_function!(
434            ConcatWsFunc::new(),
435            vec![
436                ColumnarValue::Scalar(ScalarValue::from("|")),
437                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
438            ],
439            Ok(Some("")),
440            &str,
441            Utf8,
442            StringArray
443        );
444        test_function!(
445            ConcatWsFunc::new(),
446            vec![
447                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
448                ColumnarValue::Scalar(ScalarValue::from("aa")),
449                ColumnarValue::Scalar(ScalarValue::from("bb")),
450                ColumnarValue::Scalar(ScalarValue::from("cc")),
451            ],
452            Ok(None),
453            &str,
454            Utf8,
455            StringArray
456        );
457        test_function!(
458            ConcatWsFunc::new(),
459            vec![
460                ColumnarValue::Scalar(ScalarValue::from("|")),
461                ColumnarValue::Scalar(ScalarValue::from("aa")),
462                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
463                ColumnarValue::Scalar(ScalarValue::from("cc")),
464            ],
465            Ok(Some("aa|cc")),
466            &str,
467            Utf8,
468            StringArray
469        );
470
471        Ok(())
472    }
473
474    #[test]
475    fn concat_ws() -> Result<()> {
476        // sep is scalar
477        let c0 = ColumnarValue::Scalar(ScalarValue::Utf8(Some(",".to_string())));
478        let c1 =
479            ColumnarValue::Array(Arc::new(StringArray::from(vec!["foo", "bar", "baz"])));
480        let c2 = ColumnarValue::Array(Arc::new(StringArray::from(vec![
481            Some("x"),
482            None,
483            Some("z"),
484        ])));
485
486        let arg_fields = vec![
487            Field::new("a", Utf8, true).into(),
488            Field::new("a", Utf8, true).into(),
489            Field::new("a", Utf8, true).into(),
490        ];
491        let args = ScalarFunctionArgs {
492            args: vec![c0, c1, c2],
493            arg_fields,
494            number_rows: 3,
495            return_field: Field::new("f", Utf8, true).into(),
496        };
497
498        let result = ConcatWsFunc::new().invoke_with_args(args)?;
499        let expected =
500            Arc::new(StringArray::from(vec!["foo,x", "bar", "baz,z"])) as ArrayRef;
501        match &result {
502            ColumnarValue::Array(array) => {
503                assert_eq!(&expected, array);
504            }
505            _ => panic!(),
506        }
507
508        // sep is nullable array
509        let c0 = ColumnarValue::Array(Arc::new(StringArray::from(vec![
510            Some(","),
511            None,
512            Some("+"),
513        ])));
514        let c1 =
515            ColumnarValue::Array(Arc::new(StringArray::from(vec!["foo", "bar", "baz"])));
516        let c2 = ColumnarValue::Array(Arc::new(StringArray::from(vec![
517            Some("x"),
518            Some("y"),
519            Some("z"),
520        ])));
521
522        let arg_fields = vec![
523            Field::new("a", Utf8, true).into(),
524            Field::new("a", Utf8, true).into(),
525            Field::new("a", Utf8, true).into(),
526        ];
527        let args = ScalarFunctionArgs {
528            args: vec![c0, c1, c2],
529            arg_fields,
530            number_rows: 3,
531            return_field: Field::new("f", Utf8, true).into(),
532        };
533
534        let result = ConcatWsFunc::new().invoke_with_args(args)?;
535        let expected =
536            Arc::new(StringArray::from(vec![Some("foo,x"), None, Some("baz+z")]))
537                as ArrayRef;
538        match &result {
539            ColumnarValue::Array(array) => {
540                assert_eq!(&expected, array);
541            }
542            _ => panic!(),
543        }
544
545        Ok(())
546    }
547}
datafusion_functions/string/concat_ws.rs

datafusion_functions/string/
concat_ws.rs
