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, PartialEq, Eq, Hash)]
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::config::ConfigOptions;
413    use datafusion_common::Result;
414    use datafusion_common::ScalarValue;
415    use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
416
417    use crate::utils::test::test_function;
418
419    #[test]
420    fn test_functions() -> Result<()> {
421        test_function!(
422            ConcatWsFunc::new(),
423            vec![
424                ColumnarValue::Scalar(ScalarValue::from("|")),
425                ColumnarValue::Scalar(ScalarValue::from("aa")),
426                ColumnarValue::Scalar(ScalarValue::from("bb")),
427                ColumnarValue::Scalar(ScalarValue::from("cc")),
428            ],
429            Ok(Some("aa|bb|cc")),
430            &str,
431            Utf8,
432            StringArray
433        );
434        test_function!(
435            ConcatWsFunc::new(),
436            vec![
437                ColumnarValue::Scalar(ScalarValue::from("|")),
438                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
439            ],
440            Ok(Some("")),
441            &str,
442            Utf8,
443            StringArray
444        );
445        test_function!(
446            ConcatWsFunc::new(),
447            vec![
448                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
449                ColumnarValue::Scalar(ScalarValue::from("aa")),
450                ColumnarValue::Scalar(ScalarValue::from("bb")),
451                ColumnarValue::Scalar(ScalarValue::from("cc")),
452            ],
453            Ok(None),
454            &str,
455            Utf8,
456            StringArray
457        );
458        test_function!(
459            ConcatWsFunc::new(),
460            vec![
461                ColumnarValue::Scalar(ScalarValue::from("|")),
462                ColumnarValue::Scalar(ScalarValue::from("aa")),
463                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
464                ColumnarValue::Scalar(ScalarValue::from("cc")),
465            ],
466            Ok(Some("aa|cc")),
467            &str,
468            Utf8,
469            StringArray
470        );
471
472        Ok(())
473    }
474
475    #[test]
476    fn concat_ws() -> Result<()> {
477        // sep is scalar
478        let c0 = ColumnarValue::Scalar(ScalarValue::Utf8(Some(",".to_string())));
479        let c1 =
480            ColumnarValue::Array(Arc::new(StringArray::from(vec!["foo", "bar", "baz"])));
481        let c2 = ColumnarValue::Array(Arc::new(StringArray::from(vec![
482            Some("x"),
483            None,
484            Some("z"),
485        ])));
486
487        let arg_fields = vec![
488            Field::new("a", Utf8, true).into(),
489            Field::new("a", Utf8, true).into(),
490            Field::new("a", Utf8, true).into(),
491        ];
492        let args = ScalarFunctionArgs {
493            args: vec![c0, c1, c2],
494            arg_fields,
495            number_rows: 3,
496            return_field: Field::new("f", Utf8, true).into(),
497            config_options: Arc::new(ConfigOptions::default()),
498        };
499
500        let result = ConcatWsFunc::new().invoke_with_args(args)?;
501        let expected =
502            Arc::new(StringArray::from(vec!["foo,x", "bar", "baz,z"])) as ArrayRef;
503        match &result {
504            ColumnarValue::Array(array) => {
505                assert_eq!(&expected, array);
506            }
507            _ => panic!(),
508        }
509
510        // sep is nullable array
511        let c0 = ColumnarValue::Array(Arc::new(StringArray::from(vec![
512            Some(","),
513            None,
514            Some("+"),
515        ])));
516        let c1 =
517            ColumnarValue::Array(Arc::new(StringArray::from(vec!["foo", "bar", "baz"])));
518        let c2 = ColumnarValue::Array(Arc::new(StringArray::from(vec![
519            Some("x"),
520            Some("y"),
521            Some("z"),
522        ])));
523
524        let arg_fields = vec![
525            Field::new("a", Utf8, true).into(),
526            Field::new("a", Utf8, true).into(),
527            Field::new("a", Utf8, true).into(),
528        ];
529        let args = ScalarFunctionArgs {
530            args: vec![c0, c1, c2],
531            arg_fields,
532            number_rows: 3,
533            return_field: Field::new("f", Utf8, true).into(),
534            config_options: Arc::new(ConfigOptions::default()),
535        };
536
537        let result = ConcatWsFunc::new().invoke_with_args(args)?;
538        let expected =
539            Arc::new(StringArray::from(vec![Some("foo,x"), None, Some("baz+z")]))
540                as ArrayRef;
541        match &result {
542            ColumnarValue::Array(array) => {
543                assert_eq!(&expected, array);
544            }
545            _ => panic!(),
546        }
547
548        Ok(())
549    }
550}
datafusion_functions/string/concat_ws.rs

datafusion_functions/string/
concat_ws.rs
