## 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.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};
    19
    use arrow::datatypes::DataType;
    20
    use datafusion_expr::sort_properties::ExprProperties;
    21
    use std::any::Any;
    22
    use std::sync::Arc;
    23
    24
    use crate::string::concat;
    25
    use crate::strings::{
    26
        ColumnarValueRef, LargeStringArrayBuilder, StringArrayBuilder, StringViewArrayBuilder,
    27
    };
    28
    use datafusion_common::cast::{as_string_array, as_string_view_array};
    29
    use datafusion_common::{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.",
    39
        syntax_example = "concat(str[, ..., str_n])",
    40
        sql_example = r#"```sql
    41
    > select concat('data', 'f', 'us', 'ion');
    42
    +-------------------------------------------------------+
    43
    | concat(Utf8("data"),Utf8("f"),Utf8("us"),Utf8("ion")) |
    44
    +-------------------------------------------------------+
    45
    | datafusion                                            |
    46
    +-------------------------------------------------------+
    47
    ```"#,
    48
        standard_argument(name = "str", prefix = "String"),
    49
        argument(
    50
            name = "str_n",
    51
            description = "Subsequent string expressions to concatenate."
    52
        ),
    53
        related_udf(name = "concat_ws")
    54
    )]
    55
    #[derive(Debug, PartialEq, Eq, Hash)]
    56
    pub struct ConcatFunc {
    57
        signature: Signature,
    58
    }
    59
    60
    impl Default for ConcatFunc {
    61
    0
        fn default() -> Self {
    62
    0
            ConcatFunc::new()
    63
    0
        }
    64
    }
    65
    66
    impl ConcatFunc {
    67
    0
        pub fn new() -> Self {
    68
            use DataType::*;
    69
    0
            Self {
    70
    0
                signature: Signature::variadic(
    71
    0
                    vec![Utf8View, Utf8, LargeUtf8],
    72
    0
                    Volatility::Immutable,
    73
    0
                ),
    74
    0
            }
    75
    0
        }
    76
    }
    77
    78
    impl ScalarUDFImpl for ConcatFunc {
    79
    0
        fn as_any(&self) -> &dyn Any {
    80
    0
            self
    81
    0
        }
    82
    83
    0
        fn name(&self) -> &str {
    84
    0
            "concat"
    85
    0
        }
    86
    87
    0
        fn signature(&self) -> &Signature {
    88
    0
            &self.signature
    89
    0
        }
    90
    91
    0
        fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
    92
            use DataType::*;
    93
    0
            let mut dt = &Utf8;
    94
    0
            arg_types.iter().for_each(|data_type| {
    95
    0
                if data_type == &Utf8View {
    96
    0
                    dt = data_type;
    97
    0
                }
    98
    0
                if data_type == &LargeUtf8 && dt != &Utf8View {
    99
    0
                    dt = data_type;
    100
    0
                }
    101
    0
            });
    102
    103
    0
            Ok(dt.to_owned())
    104
    0
        }
    105
    106
        /// Concatenates the text representations of all the arguments. NULL arguments are ignored.
    107
        /// concat('abcde', 2, NULL, 22) = 'abcde222'
    108
    0
        fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
    109
    0
            let ScalarFunctionArgs { args, .. } = args;
    110
    111
    0
            let mut return_datatype = DataType::Utf8;
    112
    0
            args.iter().for_each(|col| {
    113
    0
                if col.data_type() == DataType::Utf8View {
    114
    0
                    return_datatype = col.data_type();
    115
    0
                }
    116
    0
                if col.data_type() == DataType::LargeUtf8
    117
    0
                    && return_datatype != DataType::Utf8View
    118
    0
                {
    119
    0
                    return_datatype = col.data_type();
    120
    0
                }
    121
    0
            });
    122
    123
    0
            let array_len = args
    124
    0
                .iter()
    125
    0
                .filter_map(|x| match x {
    126
    0
                    ColumnarValue::Array(array) => Some(array.len()),
    127
    0
                    _ => None,
    128
    0
                })
    129
    0
                .next();
    130
    131
            // Scalar
    132
    0
            if array_len.is_none() {
    133
    0
                let mut result = String::new();
    134
    0
                for arg in args {
    135
    0
                    let ColumnarValue::Scalar(scalar) = arg else {
    136
    0
                        return internal_err!("concat expected scalar value, got {arg:?}");
    137
                    };
    138
    139
    0
                    match scalar.try_as_str() {
    140
    0
                        Some(Some(v)) => result.push_str(v),
    141
    0
                        Some(None) => {} // null literal
    142
    0
                        None => plan_err!(
    143
    0
                            "Concat function does not support scalar type {}",
    144
                            scalar
    145
    0
                        )?,
    146
                    }
    147
                }
    148
    149
    0
                return match return_datatype {
    150
                    DataType::Utf8View => {
    151
    0
                        Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(result))))
    152
                    }
    153
                    DataType::Utf8 => {
    154
    0
                        Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(result))))
    155
                    }
    156
                    DataType::LargeUtf8 => {
    157
    0
                        Ok(ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(result))))
    158
                    }
    159
    0
                    other => {
    160
    0
                        plan_err!("Concat function does not support datatype of {other}")
    161
                    }
    162
                };
    163
    0
            }
    164
    165
            // Array
    166
    0
            let len = array_len.unwrap();
    167
    0
            let mut data_size = 0;
    168
    0
            let mut columns = Vec::with_capacity(args.len());
    169
    170
    0
            for arg in &args {
    171
    0
                match arg {
    172
    0
                    ColumnarValue::Scalar(ScalarValue::Utf8(maybe_value))
    173
    0
                    | ColumnarValue::Scalar(ScalarValue::LargeUtf8(maybe_value))
    174
    0
                    | ColumnarValue::Scalar(ScalarValue::Utf8View(maybe_value)) => {
    175
    0
                        if let Some(s) = maybe_value {
    176
    0
                            data_size += s.len() * len;
    177
    0
                            columns.push(ColumnarValueRef::Scalar(s.as_bytes()));
    178
    0
                        }
    179
                    }
    180
    0
                    ColumnarValue::Array(array) => {
    181
    0
                        match array.data_type() {
    182
                            DataType::Utf8 => {
    183
    0
                                let string_array = as_string_array(array)?;
    184
    185
    0
                                data_size += string_array.values().len();
    186
    0
                                let column = if array.is_nullable() {
    187
    0
                                    ColumnarValueRef::NullableArray(string_array)
    188
                                } else {
    189
    0
                                    ColumnarValueRef::NonNullableArray(string_array)
    190
                                };
    191
    0
                                columns.push(column);
    192
                            },
    193
                            DataType::LargeUtf8 => {
    194
    0
                                let string_array = as_largestring_array(array);
    195
    196
    0
                                data_size += string_array.values().len();
    197
    0
                                let column = if array.is_nullable() {
    198
    0
                                    ColumnarValueRef::NullableLargeStringArray(string_array)
    199
                                } else {
    200
    0
                                    ColumnarValueRef::NonNullableLargeStringArray(string_array)
    201
                                };
    202
    0
                                columns.push(column);
    203
                            },
    204
                            DataType::Utf8View => {
    205
    0
                                let string_array = as_string_view_array(array)?;
    206
    207
    0
                                data_size += string_array.len();
    208
    0
                                let column = if array.is_nullable() {
    209
    0
                                    ColumnarValueRef::NullableStringViewArray(string_array)
    210
                                } else {
    211
    0
                                    ColumnarValueRef::NonNullableStringViewArray(string_array)
    212
                                };
    213
    0
                                columns.push(column);
    214
                            },
    215
    0
                            other => {
    216
    0
                                return plan_err!("Input was {other} which is not a supported datatype for concat function")
    217
                            }
    218
                        };
    219
                    }
    220
    0
                    _ => unreachable!("concat"),
    221
                }
    222
            }
    223
    224
    0
            match return_datatype {
    225
                DataType::Utf8 => {
    226
    0
                    let mut builder = StringArrayBuilder::with_capacity(len, data_size);
    227
    0
                    for i in 0..len {
    228
    0
                        columns
    229
    0
                            .iter()
    230
    0
                            .for_each(|column| builder.write::<true>(column, i));
    231
    0
                        builder.append_offset();
    232
                    }
    233
    234
    0
                    let string_array = builder.finish(None);
    235
    0
                    Ok(ColumnarValue::Array(Arc::new(string_array)))
    236
                }
    237
                DataType::Utf8View => {
    238
    0
                    let mut builder = StringViewArrayBuilder::with_capacity(len, data_size);
    239
    0
                    for i in 0..len {
    240
    0
                        columns
    241
    0
                            .iter()
    242
    0
                            .for_each(|column| builder.write::<true>(column, i));
    243
    0
                        builder.append_offset();
    244
                    }
    245
    246
    0
                    let string_array = builder.finish();
    247
    0
                    Ok(ColumnarValue::Array(Arc::new(string_array)))
    248
                }
    249
                DataType::LargeUtf8 => {
    250
    0
                    let mut builder = LargeStringArrayBuilder::with_capacity(len, data_size);
    251
    0
                    for i in 0..len {
    252
    0
                        columns
    253
    0
                            .iter()
    254
    0
                            .for_each(|column| builder.write::<true>(column, i));
    255
    0
                        builder.append_offset();
    256
                    }
    257
    258
    0
                    let string_array = builder.finish(None);
    259
    0
                    Ok(ColumnarValue::Array(Arc::new(string_array)))
    260
                }
    261
    0
                _ => unreachable!(),
    262
            }
    263
    0
        }
    264
    265
        /// Simplify the `concat` function by
    266
        /// 1. filtering out all `null` literals
    267
        /// 2. concatenating contiguous literal arguments
    268
        ///
    269
        /// For example:
    270
        /// `concat(col(a), 'hello ', 'world', col(b), null)`
    271
        /// will be optimized to
    272
        /// `concat(col(a), 'hello world', col(b))`
    273
    0
        fn simplify(
    274
    0
            &self,
    275
    0
            args: Vec<Expr>,
    276
    0
            _info: &dyn SimplifyInfo,
    277
    0
        ) -> Result<ExprSimplifyResult> {
    278
    0
            simplify_concat(args)
    279
    0
        }
    280
    281
    0
        fn documentation(&self) -> Option<&Documentation> {
    282
    0
            self.doc()
    283
    0
        }
    284
    285
    0
        fn preserves_lex_ordering(&self, _inputs: &[ExprProperties]) -> Result<bool> {
    286
    0
            Ok(true)
    287
    0
        }
    288
    }
    289
    290
    0
    pub fn simplify_concat(args: Vec<Expr>) -> Result<ExprSimplifyResult> {
    291
    0
        let mut new_args = Vec::with_capacity(args.len());
    292
    0
        let mut contiguous_scalar = "".to_string();
    293
    294
    0
        let return_type = {
    295
    0
            let data_types: Vec<_> = args
    296
    0
                .iter()
    297
    0
                .filter_map(|expr| match expr {
    298
    0
                    Expr::Literal(l, _) => Some(l.data_type()),
    299
    0
                    _ => None,
    300
    0
                })
    301
    0
                .collect();
    302
    0
            ConcatFunc::new().return_type(&data_types)
    303
    0
        }?;
    304
    305
    0
        for arg in args.clone() {
    306
    0
            match arg {
    307
    0
                Expr::Literal(ScalarValue::Utf8(None), _) => {}
    308
    0
                Expr::Literal(ScalarValue::LargeUtf8(None), _) => {
    309
    0
                }
    310
    0
                Expr::Literal(ScalarValue::Utf8View(None), _) => { }
    311
    312
                // filter out `null` args
    313
                // All literals have been converted to Utf8 or LargeUtf8 in type_coercion.
    314
                // Concatenate it with the `contiguous_scalar`.
    315
    0
                Expr::Literal(ScalarValue::Utf8(Some(v)), _) => {
    316
    0
                    contiguous_scalar += &v;
    317
    0
                }
    318
    0
                Expr::Literal(ScalarValue::LargeUtf8(Some(v)), _) => {
    319
    0
                    contiguous_scalar += &v;
    320
    0
                }
    321
    0
                Expr::Literal(ScalarValue::Utf8View(Some(v)), _) => {
    322
    0
                    contiguous_scalar += &v;
    323
    0
                }
    324
    325
    0
                Expr::Literal(x, _) => {
    326
    0
                    return internal_err!(
    327
    0
                        "The scalar {x} should be casted to string type during the type coercion."
    328
                    )
    329
                }
    330
                // If the arg is not a literal, we should first push the current `contiguous_scalar`
    331
                // to the `new_args` (if it is not empty) and reset it to empty string.
    332
                // Then pushing this arg to the `new_args`.
    333
    0
                arg => {
    334
    0
                    if !contiguous_scalar.is_empty() {
    335
    0
                        match return_type {
    336
    0
                            DataType::Utf8 => new_args.push(lit(contiguous_scalar)),
    337
    0
                            DataType::LargeUtf8 => new_args.push(lit(ScalarValue::LargeUtf8(Some(contiguous_scalar)))),
    338
    0
                            DataType::Utf8View => new_args.push(lit(ScalarValue::Utf8View(Some(contiguous_scalar)))),
    339
    0
                            _ => unreachable!(),
    340
                        }
    341
    0
                        contiguous_scalar = "".to_string();
    342
    0
                    }
    343
    0
                    new_args.push(arg);
    344
                }
    345
            }
    346
        }
    347
    348
    0
        if !contiguous_scalar.is_empty() {
    349
    0
            match return_type {
    350
    0
                DataType::Utf8 => new_args.push(lit(contiguous_scalar)),
    351
                DataType::LargeUtf8 => {
    352
    0
                    new_args.push(lit(ScalarValue::LargeUtf8(Some(contiguous_scalar))))
    353
                }
    354
                DataType::Utf8View => {
    355
    0
                    new_args.push(lit(ScalarValue::Utf8View(Some(contiguous_scalar))))
    356
                }
    357
    0
                _ => unreachable!(),
    358
            }
    359
    0
        }
    360
    361
    0
        if !args.eq(&new_args) {
    362
    0
            Ok(ExprSimplifyResult::Simplified(Expr::ScalarFunction(
    363
    0
                ScalarFunction {
    364
    0
                    func: concat(),
    365
    0
                    args: new_args,
    366
    0
                },
    367
    0
            )))
    368
        } else {
    369
    0
            Ok(ExprSimplifyResult::Original(args))
    370
        }
    371
    0
    }
    372
    373
    #[cfg(test)]
    374
    mod tests {
    375
        use super::*;
    376
        use crate::utils::test::test_function;
    377
        use arrow::array::{Array, LargeStringArray, StringViewArray};
    378
        use arrow::array::{ArrayRef, StringArray};
    379
        use arrow::datatypes::Field;
    380
        use datafusion_common::config::ConfigOptions;
    381
        use DataType::*;
    382
    383
        #[test]
    384
        fn test_functions() -> Result<()> {
    385
            test_function!(
    386
                ConcatFunc::new(),
    387
                vec![
    388
                    ColumnarValue::Scalar(ScalarValue::from("aa")),
    389
                    ColumnarValue::Scalar(ScalarValue::from("bb")),
    390
                    ColumnarValue::Scalar(ScalarValue::from("cc")),
    391
                ],
    392
                Ok(Some("aabbcc")),
    393
                &str,
    394
                Utf8,
    395
                StringArray
    396
            );
    397
            test_function!(
    398
                ConcatFunc::new(),
    399
                vec![
    400
                    ColumnarValue::Scalar(ScalarValue::from("aa")),
    401
                    ColumnarValue::Scalar(ScalarValue::Utf8(None)),
    402
                    ColumnarValue::Scalar(ScalarValue::from("cc")),
    403
                ],
    404
                Ok(Some("aacc")),
    405
                &str,
    406
                Utf8,
    407
                StringArray
    408
            );
    409
            test_function!(
    410
                ConcatFunc::new(),
    411
                vec![ColumnarValue::Scalar(ScalarValue::Utf8(None))],
    412
                Ok(Some("")),
    413
                &str,
    414
                Utf8,
    415
                StringArray
    416
            );
    417
            test_function!(
    418
                ConcatFunc::new(),
    419
                vec![
    420
                    ColumnarValue::Scalar(ScalarValue::from("aa")),
    421
                    ColumnarValue::Scalar(ScalarValue::Utf8View(None)),
    422
                    ColumnarValue::Scalar(ScalarValue::LargeUtf8(None)),
    423
                    ColumnarValue::Scalar(ScalarValue::from("cc")),
    424
                ],
    425
                Ok(Some("aacc")),
    426
                &str,
    427
                Utf8View,
    428
                StringViewArray
    429
            );
    430
            test_function!(
    431
                ConcatFunc::new(),
    432
                vec![
    433
                    ColumnarValue::Scalar(ScalarValue::from("aa")),
    434
                    ColumnarValue::Scalar(ScalarValue::LargeUtf8(None)),
    435
                    ColumnarValue::Scalar(ScalarValue::from("cc")),
    436
                ],
    437
                Ok(Some("aacc")),
    438
                &str,
    439
                LargeUtf8,
    440
                LargeStringArray
    441
            );
    442
            test_function!(
    443
                ConcatFunc::new(),
    444
                vec![
    445
                    ColumnarValue::Scalar(ScalarValue::Utf8View(Some("aa".to_string()))),
    446
                    ColumnarValue::Scalar(ScalarValue::Utf8(Some("cc".to_string()))),
    447
                ],
    448
                Ok(Some("aacc")),
    449
                &str,
    450
                Utf8View,
    451
                StringViewArray
    452
            );
    453
    454
            Ok(())
    455
        }
    456
    457
        #[test]
    458
        fn concat() -> Result<()> {
    459
            let c0 =
    460
                ColumnarValue::Array(Arc::new(StringArray::from(vec!["foo", "bar", "baz"])));
    461
            let c1 = ColumnarValue::Scalar(ScalarValue::Utf8(Some(",".to_string())));
    462
            let c2 = ColumnarValue::Array(Arc::new(StringArray::from(vec![
    463
                Some("x"),
    464
                None,
    465
                Some("z"),
    466
            ])));
    467
            let c3 = ColumnarValue::Scalar(ScalarValue::Utf8View(Some(",".to_string())));
    468
            let c4 = ColumnarValue::Array(Arc::new(StringViewArray::from(vec![
    469
                Some("a"),
    470
                None,
    471
                Some("b"),
    472
            ])));
    473
            let arg_fields = vec![
    474
                Field::new("a", Utf8, true),
    475
                Field::new("a", Utf8, true),
    476
                Field::new("a", Utf8, true),
    477
                Field::new("a", Utf8View, true),
    478
                Field::new("a", Utf8View, true),
    479
            ]
    480
            .into_iter()
    481
            .map(Arc::new)
    482
            .collect::<Vec<_>>();
    483
    484
            let args = ScalarFunctionArgs {
    485
                args: vec![c0, c1, c2, c3, c4],
    486
                arg_fields,
    487
                number_rows: 3,
    488
                return_field: Field::new("f", Utf8, true).into(),
    489
                config_options: Arc::new(ConfigOptions::default()),
    490
            };
    491
    492
            let result = ConcatFunc::new().invoke_with_args(args)?;
    493
            let expected =
    494
                Arc::new(StringViewArray::from(vec!["foo,x,a", "bar,,", "baz,z,b"]))
    495
                    as ArrayRef;
    496
            match &result {
    497
                ColumnarValue::Array(array) => {
    498
                    assert_eq!(&expected, array);
    499
                }
    500
                _ => panic!(),
    501
            }
    502
            Ok(())
    503
        }
    504
    }