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            match delimiter {
317                // when the delimiter is an empty string,
318                // we can use `concat` to replace `concat_ws`
319                Some(delimiter) if delimiter.is_empty() => {
320                    match simplify_concat(args.to_vec())? {
321                        ExprSimplifyResult::Original(_) => {
322                            Ok(ExprSimplifyResult::Simplified(Expr::ScalarFunction(
323                                ScalarFunction {
324                                    func: concat(),
325                                    args: args.to_vec(),
326                                },
327                            )))
328                        }
329                        expr => Ok(expr),
330                    }
331                }
332                Some(delimiter) => {
333                    let mut new_args = Vec::with_capacity(args.len());
334                    new_args.push(lit(delimiter));
335                    let mut contiguous_scalar = None;
336                    for arg in args {
337                        match arg {
338                            // filter out null args
339                            Expr::Literal(ScalarValue::Utf8(None) | ScalarValue::LargeUtf8(None) | ScalarValue::Utf8View(None)) => {}
340                            Expr::Literal(ScalarValue::Utf8(Some(v)) | ScalarValue::LargeUtf8(Some(v)) | ScalarValue::Utf8View(Some(v))) => {
341                                match contiguous_scalar {
342                                    None => contiguous_scalar = Some(v.to_string()),
343                                    Some(mut pre) => {
344                                        pre += delimiter;
345                                        pre += v;
346                                        contiguous_scalar = Some(pre)
347                                    }
348                                }
349                            }
350                            Expr::Literal(s) => return internal_err!("The scalar {s} should be casted to string type during the type coercion."),
351                            // If the arg is not a literal, we should first push the current `contiguous_scalar`
352                            // to the `new_args` and reset it to None.
353                            // Then pushing this arg to the `new_args`.
354                            arg => {
355                                if let Some(val) = contiguous_scalar {
356                                    new_args.push(lit(val));
357                                }
358                                new_args.push(arg.clone());
359                                contiguous_scalar = None;
360                            }
361                        }
362                    }
363                    if let Some(val) = contiguous_scalar {
364                        new_args.push(lit(val));
365                    }
366
367                    Ok(ExprSimplifyResult::Simplified(Expr::ScalarFunction(
368                        ScalarFunction {
369                            func: concat_ws(),
370                            args: new_args,
371                        },
372                    )))
373                }
374                // if the delimiter is null, then the value of the whole expression is null.
375                None => Ok(ExprSimplifyResult::Simplified(Expr::Literal(
376                    ScalarValue::Utf8(None),
377                ))),
378            }
379        }
380        Expr::Literal(d) => internal_err!(
381            "The scalar {d} should be casted to string type during the type coercion."
382        ),
383        _ => {
384            let mut args = args
385                .iter()
386                .filter(|&x| !is_null(x))
387                .cloned()
388                .collect::<Vec<Expr>>();
389            args.insert(0, delimiter.clone());
390            Ok(ExprSimplifyResult::Original(args))
391        }
392    }
393}
394
395fn is_null(expr: &Expr) -> bool {
396    match expr {
397        Expr::Literal(v) => v.is_null(),
398        _ => false,
399    }
400}
401
402#[cfg(test)]
403mod tests {
404    use std::sync::Arc;
405
406    use arrow::array::{Array, ArrayRef, StringArray};
407    use arrow::datatypes::DataType::Utf8;
408
409    use crate::string::concat_ws::ConcatWsFunc;
410    use datafusion_common::Result;
411    use datafusion_common::ScalarValue;
412    use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
413
414    use crate::utils::test::test_function;
415
416    #[test]
417    fn test_functions() -> Result<()> {
418        test_function!(
419            ConcatWsFunc::new(),
420            vec![
421                ColumnarValue::Scalar(ScalarValue::from("|")),
422                ColumnarValue::Scalar(ScalarValue::from("aa")),
423                ColumnarValue::Scalar(ScalarValue::from("bb")),
424                ColumnarValue::Scalar(ScalarValue::from("cc")),
425            ],
426            Ok(Some("aa|bb|cc")),
427            &str,
428            Utf8,
429            StringArray
430        );
431        test_function!(
432            ConcatWsFunc::new(),
433            vec![
434                ColumnarValue::Scalar(ScalarValue::from("|")),
435                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
436            ],
437            Ok(Some("")),
438            &str,
439            Utf8,
440            StringArray
441        );
442        test_function!(
443            ConcatWsFunc::new(),
444            vec![
445                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
446                ColumnarValue::Scalar(ScalarValue::from("aa")),
447                ColumnarValue::Scalar(ScalarValue::from("bb")),
448                ColumnarValue::Scalar(ScalarValue::from("cc")),
449            ],
450            Ok(None),
451            &str,
452            Utf8,
453            StringArray
454        );
455        test_function!(
456            ConcatWsFunc::new(),
457            vec![
458                ColumnarValue::Scalar(ScalarValue::from("|")),
459                ColumnarValue::Scalar(ScalarValue::from("aa")),
460                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
461                ColumnarValue::Scalar(ScalarValue::from("cc")),
462            ],
463            Ok(Some("aa|cc")),
464            &str,
465            Utf8,
466            StringArray
467        );
468
469        Ok(())
470    }
471
472    #[test]
473    fn concat_ws() -> Result<()> {
474        // sep is scalar
475        let c0 = ColumnarValue::Scalar(ScalarValue::Utf8(Some(",".to_string())));
476        let c1 =
477            ColumnarValue::Array(Arc::new(StringArray::from(vec!["foo", "bar", "baz"])));
478        let c2 = ColumnarValue::Array(Arc::new(StringArray::from(vec![
479            Some("x"),
480            None,
481            Some("z"),
482        ])));
483
484        let args = ScalarFunctionArgs {
485            args: vec![c0, c1, c2],
486            number_rows: 3,
487            return_type: &Utf8,
488        };
489
490        let result = ConcatWsFunc::new().invoke_with_args(args)?;
491        let expected =
492            Arc::new(StringArray::from(vec!["foo,x", "bar", "baz,z"])) as ArrayRef;
493        match &result {
494            ColumnarValue::Array(array) => {
495                assert_eq!(&expected, array);
496            }
497            _ => panic!(),
498        }
499
500        // sep is nullable array
501        let c0 = ColumnarValue::Array(Arc::new(StringArray::from(vec![
502            Some(","),
503            None,
504            Some("+"),
505        ])));
506        let c1 =
507            ColumnarValue::Array(Arc::new(StringArray::from(vec!["foo", "bar", "baz"])));
508        let c2 = ColumnarValue::Array(Arc::new(StringArray::from(vec![
509            Some("x"),
510            Some("y"),
511            Some("z"),
512        ])));
513
514        let args = ScalarFunctionArgs {
515            args: vec![c0, c1, c2],
516            number_rows: 3,
517            return_type: &Utf8,
518        };
519
520        let result = ConcatWsFunc::new().invoke_with_args(args)?;
521        let expected =
522            Arc::new(StringArray::from(vec![Some("foo,x"), None, Some("baz+z")]))
523                as ArrayRef;
524        match &result {
525            ColumnarValue::Array(array) => {
526                assert_eq!(&expected, array);
527            }
528            _ => panic!(),
529        }
530
531        Ok(())
532    }
533}
datafusion_functions/string/concat_ws.rs

datafusion_functions/string/
concat_ws.rs
