## Coverage Report

#### Created: 2025-11-22 06:00

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /Users/andrewlamb/Software/datafusion/datafusion/functions/src/string/concat_ws.rs

    Line
    Count
    Source
    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
    use arrow::array::{as_largestring_array, Array, StringArray};
    19
    use std::any::Any;
    20
    use std::sync::Arc;
    21
    22
    use arrow::datatypes::DataType;
    23
    24
    use crate::string::concat;
    25
    use crate::string::concat::simplify_concat;
    26
    use crate::string::concat_ws;
    27
    use crate::strings::{ColumnarValueRef, StringArrayBuilder};
    28
    use datafusion_common::cast::{as_string_array, as_string_view_array};
    29
    use datafusion_common::{exec_err, internal_err, plan_err, Result, ScalarValue};
    30
    use datafusion_expr::expr::ScalarFunction;
    31
    use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
    32
    use datafusion_expr::{lit, ColumnarValue, Documentation, Expr, Volatility};
    33
    use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature};
    34
    use 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)]
    63
    pub struct ConcatWsFunc {
    64
        signature: Signature,
    65
    }
    66
    67
    impl Default for ConcatWsFunc {
    68
    0
        fn default() -> Self {
    69
    0
            ConcatWsFunc::new()
    70
    0
        }
    71
    }
    72
    73
    impl ConcatWsFunc {
    74
    0
        pub fn new() -> Self {
    75
            use DataType::*;
    76
    0
            Self {
    77
    0
                signature: Signature::variadic(
    78
    0
                    vec![Utf8View, Utf8, LargeUtf8],
    79
    0
                    Volatility::Immutable,
    80
    0
                ),
    81
    0
            }
    82
    0
        }
    83
    }
    84
    85
    impl ScalarUDFImpl for ConcatWsFunc {
    86
    0
        fn as_any(&self) -> &dyn Any {
    87
    0
            self
    88
    0
        }
    89
    90
    0
        fn name(&self) -> &str {
    91
    0
            "concat_ws"
    92
    0
        }
    93
    94
    0
        fn signature(&self) -> &Signature {
    95
    0
            &self.signature
    96
    0
        }
    97
    98
    0
        fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
    99
            use DataType::*;
    100
    0
            Ok(Utf8)
    101
    0
        }
    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
    0
        fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
    106
    0
            let ScalarFunctionArgs { args, .. } = args;
    107
    108
            // do not accept 0 arguments.
    109
    0
            if args.len() < 2 {
    110
    0
                return exec_err!(
    111
    0
                    "concat_ws was called with {} arguments. It requires at least 2.",
    112
    0
                    args.len()
    113
                );
    114
    0
            }
    115
    116
    0
            let array_len = args
    117
    0
                .iter()
    118
    0
                .filter_map(|x| match x {
    119
    0
                    ColumnarValue::Array(array) => Some(array.len()),
    120
    0
                    _ => None,
    121
    0
                })
    122
    0
                .next();
    123
    124
            // Scalar
    125
    0
            if array_len.is_none() {
    126
    0
                let ColumnarValue::Scalar(scalar) = &args[0] else {
    127
                    // loop above checks for all args being scalar
    128
    0
                    unreachable!()
    129
                };
    130
    0
                let sep = match scalar.try_as_str() {
    131
    0
                    Some(Some(s)) => s,
    132
                    Some(None) => {
    133
                        // null literal string
    134
    0
                        return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None)));
    135
                    }
    136
    0
                    None => return internal_err!("Expected string literal, got {scalar:?}"),
    137
                };
    138
    139
    0
                let mut result = String::new();
    140
                // iterator over Option<str>
    141
    0
                let iter = &mut args[1..].iter().map(|arg| {
    142
    0
                    let ColumnarValue::Scalar(scalar) = arg else {
    143
                        // loop above checks for all args being scalar
    144
    0
                        unreachable!()
    145
                    };
    146
    0
                    scalar.try_as_str()
    147
    0
                });
    148
    149
                // append first non null arg
    150
    0
                for scalar in iter.by_ref() {
    151
    0
                    match scalar {
    152
    0
                        Some(Some(s)) => {
    153
    0
                            result.push_str(s);
    154
    0
                            break;
    155
                        }
    156
    0
                        Some(None) => {} // null literal string
    157
                        None => {
    158
    0
                            return internal_err!("Expected string literal, got {scalar:?}")
    159
                        }
    160
                    }
    161
                }
    162
    163
                // handle subsequent non null args
    164
    0
                for scalar in iter.by_ref() {
    165
    0
                    match scalar {
    166
    0
                        Some(Some(s)) => {
    167
    0
                            result.push_str(sep);
    168
    0
                            result.push_str(s);
    169
    0
                        }
    170
    0
                        Some(None) => {} // null literal string
    171
                        None => {
    172
    0
                            return internal_err!("Expected string literal, got {scalar:?}")
    173
                        }
    174
                    }
    175
                }
    176
    177
    0
                return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(result))));
    178
    0
            }
    179
    180
            // Array
    181
    0
            let len = array_len.unwrap();
    182
    0
            let mut data_size = 0;
    183
    184
            // parse sep
    185
    0
            let sep = match &args[0] {
    186
    0
                ColumnarValue::Scalar(ScalarValue::Utf8(Some(s))) => {
    187
    0
                    data_size += s.len() * len * (args.len() - 2); // estimate
    188
    0
                    ColumnarValueRef::Scalar(s.as_bytes())
    189
                }
    190
                ColumnarValue::Scalar(ScalarValue::Utf8(None)) => {
    191
    0
                    return Ok(ColumnarValue::Array(Arc::new(StringArray::new_null(len))));
    192
                }
    193
    0
                ColumnarValue::Array(array) => {
    194
    0
                    let string_array = as_string_array(array)?;
    195
    0
                    data_size += string_array.values().len() * (args.len() - 2); // estimate
    196
    0
                    if array.is_nullable() {
    197
    0
                        ColumnarValueRef::NullableArray(string_array)
    198
                    } else {
    199
    0
                        ColumnarValueRef::NonNullableArray(string_array)
    200
                    }
    201
                }
    202
    0
                _ => unreachable!("concat ws"),
    203
            };
    204
    205
    0
            let mut columns = Vec::with_capacity(args.len() - 1);
    206
    0
            for arg in &args[1..] {
    207
    0
                match arg {
    208
    0
                    ColumnarValue::Scalar(ScalarValue::Utf8(maybe_value))
    209
    0
                    | ColumnarValue::Scalar(ScalarValue::LargeUtf8(maybe_value))
    210
    0
                    | ColumnarValue::Scalar(ScalarValue::Utf8View(maybe_value)) => {
    211
    0
                        if let Some(s) = maybe_value {
    212
    0
                            data_size += s.len() * len;
    213
    0
                            columns.push(ColumnarValueRef::Scalar(s.as_bytes()));
    214
    0
                        }
    215
                    }
    216
    0
                    ColumnarValue::Array(array) => {
    217
    0
                        match array.data_type() {
    218
                            DataType::Utf8 => {
    219
    0
                                let string_array = as_string_array(array)?;
    220
    221
    0
                                data_size += string_array.values().len();
    222
    0
                                let column = if array.is_nullable() {
    223
    0
                                    ColumnarValueRef::NullableArray(string_array)
    224
                                } else {
    225
    0
                                    ColumnarValueRef::NonNullableArray(string_array)
    226
                                };
    227
    0
                                columns.push(column);
    228
                            },
    229
                            DataType::LargeUtf8 => {
    230
    0
                                let string_array = as_largestring_array(array);
    231
    232
    0
                                data_size += string_array.values().len();
    233
    0
                                let column = if array.is_nullable() {
    234
    0
                                    ColumnarValueRef::NullableLargeStringArray(string_array)
    235
                                } else {
    236
    0
                                    ColumnarValueRef::NonNullableLargeStringArray(string_array)
    237
                                };
    238
    0
                                columns.push(column);
    239
                            },
    240
                            DataType::Utf8View => {
    241
    0
                                let string_array = as_string_view_array(array)?;
    242
    243
    0
                                data_size += string_array.data_buffers().iter().map(|buf| buf.len()).sum::<usize>();
    244
    0
                                let column = if array.is_nullable() {
    245
    0
                                    ColumnarValueRef::NullableStringViewArray(string_array)
    246
                                } else {
    247
    0
                                    ColumnarValueRef::NonNullableStringViewArray(string_array)
    248
                                };
    249
    0
                                columns.push(column);
    250
                            },
    251
    0
                            other => {
    252
    0
                                return plan_err!("Input was {other} which is not a supported datatype for concat_ws function.")
    253
                            }
    254
                        };
    255
                    }
    256
    0
                    _ => unreachable!(),
    257
                }
    258
            }
    259
    260
    0
            let mut builder = StringArrayBuilder::with_capacity(len, data_size);
    261
    0
            for i in 0..len {
    262
    0
                if !sep.is_valid(i) {
    263
    0
                    builder.append_offset();
    264
    0
                    continue;
    265
    0
                }
    266
    267
    0
                let mut iter = columns.iter();
    268
    0
                for column in iter.by_ref() {
    269
    0
                    if column.is_valid(i) {
    270
    0
                        builder.write::<false>(column, i);
    271
    0
                        break;
    272
    0
                    }
    273
                }
    274
    275
    0
                for column in iter {
    276
    0
                    if column.is_valid(i) {
    277
    0
                        builder.write::<false>(&sep, i);
    278
    0
                        builder.write::<false>(column, i);
    279
    0
                    }
    280
                }
    281
    282
    0
                builder.append_offset();
    283
            }
    284
    285
    0
            Ok(ColumnarValue::Array(Arc::new(builder.finish(sep.nulls()))))
    286
    0
        }
    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
    0
        fn simplify(
    294
    0
            &self,
    295
    0
            args: Vec<Expr>,
    296
    0
            _info: &dyn SimplifyInfo,
    297
    0
        ) -> Result<ExprSimplifyResult> {
    298
    0
            match &args[..] {
    299
    0
                [delimiter, vals @ ..] => simplify_concat_ws(delimiter, vals),
    300
    0
                _ => Ok(ExprSimplifyResult::Original(args)),
    301
            }
    302
    0
        }
    303
    304
    0
        fn documentation(&self) -> Option<&Documentation> {
    305
    0
            self.doc()
    306
    0
        }
    307
    }
    308
    309
    0
    fn simplify_concat_ws(delimiter: &Expr, args: &[Expr]) -> Result<ExprSimplifyResult> {
    310
    0
        match delimiter {
    311
            Expr::Literal(
    312
    0
                ScalarValue::Utf8(delimiter)
    313
    0
                | ScalarValue::LargeUtf8(delimiter)
    314
    0
                | ScalarValue::Utf8View(delimiter),
    315
                _,
    316
            ) => {
    317
    0
                match delimiter {
    318
                    // when the delimiter is an empty string,
    319
                    // we can use `concat` to replace `concat_ws`
    320
    0
                    Some(delimiter) if delimiter.is_empty() => {
    321
    0
                        match simplify_concat(args.to_vec())? {
    322
                            ExprSimplifyResult::Original(_) => {
    323
    0
                                Ok(ExprSimplifyResult::Simplified(Expr::ScalarFunction(
    324
    0
                                    ScalarFunction {
    325
    0
                                        func: concat(),
    326
    0
                                        args: args.to_vec(),
    327
    0
                                    },
    328
    0
                                )))
    329
                            }
    330
    0
                            expr => Ok(expr),
    331
                        }
    332
                    }
    333
    0
                    Some(delimiter) => {
    334
    0
                        let mut new_args = Vec::with_capacity(args.len());
    335
    0
                        new_args.push(lit(delimiter));
    336
    0
                        let mut contiguous_scalar = None;
    337
    0
                        for arg in args {
    338
    0
                            match arg {
    339
                                // filter out null args
    340
    0
                                Expr::Literal(ScalarValue::Utf8(None) | ScalarValue::LargeUtf8(None) | ScalarValue::Utf8View(None), _) => {}
    341
    0
                                Expr::Literal(ScalarValue::Utf8(Some(v)) | ScalarValue::LargeUtf8(Some(v)) | ScalarValue::Utf8View(Some(v)), _) => {
    342
    0
                                    match contiguous_scalar {
    343
    0
                                        None => contiguous_scalar = Some(v.to_string()),
    344
    0
                                        Some(mut pre) => {
    345
    0
                                            pre += delimiter;
    346
    0
                                            pre += v;
    347
    0
                                            contiguous_scalar = Some(pre)
    348
                                        }
    349
                                    }
    350
                                }
    351
    0
                                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
    0
                                arg => {
    356
    0
                                    if let Some(val) = contiguous_scalar {
    357
    0
                                        new_args.push(lit(val));
    358
    0
                                    }
    359
    0
                                    new_args.push(arg.clone());
    360
    0
                                    contiguous_scalar = None;
    361
                                }
    362
                            }
    363
                        }
    364
    0
                        if let Some(val) = contiguous_scalar {
    365
    0
                            new_args.push(lit(val));
    366
    0
                        }
    367
    368
    0
                        Ok(ExprSimplifyResult::Simplified(Expr::ScalarFunction(
    369
    0
                            ScalarFunction {
    370
    0
                                func: concat_ws(),
    371
    0
                                args: new_args,
    372
    0
                            },
    373
    0
                        )))
    374
                    }
    375
                    // if the delimiter is null, then the value of the whole expression is null.
    376
    0
                    None => Ok(ExprSimplifyResult::Simplified(Expr::Literal(
    377
    0
                        ScalarValue::Utf8(None),
    378
    0
                        None,
    379
    0
                    ))),
    380
                }
    381
            }
    382
    0
            Expr::Literal(d, _) => internal_err!(
    383
    0
                "The scalar {d} should be casted to string type during the type coercion."
    384
            ),
    385
            _ => {
    386
    0
                let mut args = args
    387
    0
                    .iter()
    388
    0
                    .filter(|&x| !is_null(x))
    389
    0
                    .cloned()
    390
    0
                    .collect::<Vec<Expr>>();
    391
    0
                args.insert(0, delimiter.clone());
    392
    0
                Ok(ExprSimplifyResult::Original(args))
    393
            }
    394
        }
    395
    0
    }
    396
    397
    0
    fn is_null(expr: &Expr) -> bool {
    398
    0
        match expr {
    399
    0
            Expr::Literal(v, _) => v.is_null(),
    400
    0
            _ => false,
    401
        }
    402
    0
    }
    403
    404
    #[cfg(test)]
    405
    mod 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
    }