## 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/split_part.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 crate::utils::utf8_to_str_type;
    19
    use arrow::array::{
    20
        ArrayRef, GenericStringArray, Int64Array, OffsetSizeTrait, StringArrayType,
    21
        StringViewArray,
    22
    };
    23
    use arrow::array::{AsArray, GenericStringBuilder};
    24
    use arrow::datatypes::DataType;
    25
    use datafusion_common::cast::as_int64_array;
    26
    use datafusion_common::ScalarValue;
    27
    use datafusion_common::{exec_err, DataFusionError, Result};
    28
    use datafusion_expr::{ColumnarValue, Documentation, TypeSignature, Volatility};
    29
    use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature};
    30
    use datafusion_macros::user_doc;
    31
    use std::any::Any;
    32
    use std::sync::Arc;
    33
    34
    #[user_doc(
    35
        doc_section(label = "String Functions"),
    36
        description = "Splits a string based on a specified delimiter and returns the substring in the specified position.",
    37
        syntax_example = "split_part(str, delimiter, pos)",
    38
        sql_example = r#"```sql
    39
    > select split_part('1.2.3.4.5', '.', 3);
    40
    +--------------------------------------------------+
    41
    | split_part(Utf8("1.2.3.4.5"),Utf8("."),Int64(3)) |
    42
    +--------------------------------------------------+
    43
    | 3                                                |
    44
    +--------------------------------------------------+
    45
    ```"#,
    46
        standard_argument(name = "str", prefix = "String"),
    47
        argument(name = "delimiter", description = "String or character to split on."),
    48
        argument(name = "pos", description = "Position of the part to return.")
    49
    )]
    50
    #[derive(Debug, PartialEq, Eq, Hash)]
    51
    pub struct SplitPartFunc {
    52
        signature: Signature,
    53
    }
    54
    55
    impl Default for SplitPartFunc {
    56
    0
        fn default() -> Self {
    57
    0
            Self::new()
    58
    0
        }
    59
    }
    60
    61
    impl SplitPartFunc {
    62
    0
        pub fn new() -> Self {
    63
            use DataType::*;
    64
    0
            Self {
    65
    0
                signature: Signature::one_of(
    66
    0
                    vec![
    67
    0
                        TypeSignature::Exact(vec![Utf8View, Utf8View, Int64]),
    68
    0
                        TypeSignature::Exact(vec![Utf8View, Utf8, Int64]),
    69
    0
                        TypeSignature::Exact(vec![Utf8View, LargeUtf8, Int64]),
    70
    0
                        TypeSignature::Exact(vec![Utf8, Utf8View, Int64]),
    71
    0
                        TypeSignature::Exact(vec![Utf8, Utf8, Int64]),
    72
    0
                        TypeSignature::Exact(vec![LargeUtf8, Utf8View, Int64]),
    73
    0
                        TypeSignature::Exact(vec![LargeUtf8, Utf8, Int64]),
    74
    0
                        TypeSignature::Exact(vec![Utf8, LargeUtf8, Int64]),
    75
    0
                        TypeSignature::Exact(vec![LargeUtf8, LargeUtf8, Int64]),
    76
    0
                    ],
    77
    0
                    Volatility::Immutable,
    78
    0
                ),
    79
    0
            }
    80
    0
        }
    81
    }
    82
    83
    impl ScalarUDFImpl for SplitPartFunc {
    84
    0
        fn as_any(&self) -> &dyn Any {
    85
    0
            self
    86
    0
        }
    87
    88
    0
        fn name(&self) -> &str {
    89
    0
            "split_part"
    90
    0
        }
    91
    92
    0
        fn signature(&self) -> &Signature {
    93
    0
            &self.signature
    94
    0
        }
    95
    96
    0
        fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
    97
    0
            utf8_to_str_type(&arg_types[0], "split_part")
    98
    0
        }
    99
    100
    0
        fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
    101
    0
            let ScalarFunctionArgs { args, .. } = args;
    102
    103
            // First, determine if any of the arguments is an Array
    104
    0
            let len = args.iter().find_map(|arg| match arg {
    105
    0
                ColumnarValue::Array(a) => Some(a.len()),
    106
    0
                _ => None,
    107
    0
            });
    108
    109
    0
            let inferred_length = len.unwrap_or(1);
    110
    0
            let is_scalar = len.is_none();
    111
    112
            // Convert all ColumnarValues to ArrayRefs
    113
    0
            let args = args
    114
    0
                .iter()
    115
    0
                .map(|arg| match arg {
    116
    0
                    ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(inferred_length),
    117
    0
                    ColumnarValue::Array(array) => Ok(Arc::clone(array)),
    118
    0
                })
    119
    0
                .collect::<Result<Vec<_>>>()?;
    120
    121
            // Unpack the ArrayRefs from the arguments
    122
    0
            let n_array = as_int64_array(&args[2])?;
    123
    0
            let result = match (args[0].data_type(), args[1].data_type()) {
    124
                (DataType::Utf8View, DataType::Utf8View) => {
    125
    0
                    split_part_impl::<&StringViewArray, &StringViewArray, i32>(
    126
    0
                        args[0].as_string_view(),
    127
    0
                        args[1].as_string_view(),
    128
    0
                        n_array,
    129
                    )
    130
                }
    131
                (DataType::Utf8View, DataType::Utf8) => {
    132
    0
                    split_part_impl::<&StringViewArray, &GenericStringArray<i32>, i32>(
    133
    0
                        args[0].as_string_view(),
    134
    0
                        args[1].as_string::<i32>(),
    135
    0
                        n_array,
    136
                    )
    137
                }
    138
                (DataType::Utf8View, DataType::LargeUtf8) => {
    139
    0
                    split_part_impl::<&StringViewArray, &GenericStringArray<i64>, i32>(
    140
    0
                        args[0].as_string_view(),
    141
    0
                        args[1].as_string::<i64>(),
    142
    0
                        n_array,
    143
                    )
    144
                }
    145
                (DataType::Utf8, DataType::Utf8View) => {
    146
    0
                    split_part_impl::<&GenericStringArray<i32>, &StringViewArray, i32>(
    147
    0
                        args[0].as_string::<i32>(),
    148
    0
                        args[1].as_string_view(),
    149
    0
                        n_array,
    150
                    )
    151
                }
    152
                (DataType::LargeUtf8, DataType::Utf8View) => {
    153
    0
                    split_part_impl::<&GenericStringArray<i64>, &StringViewArray, i64>(
    154
    0
                        args[0].as_string::<i64>(),
    155
    0
                        args[1].as_string_view(),
    156
    0
                        n_array,
    157
                    )
    158
                }
    159
                (DataType::Utf8, DataType::Utf8) => {
    160
    0
                    split_part_impl::<&GenericStringArray<i32>, &GenericStringArray<i32>, i32>(
    161
    0
                        args[0].as_string::<i32>(),
    162
    0
                        args[1].as_string::<i32>(),
    163
    0
                        n_array,
    164
                    )
    165
                }
    166
                (DataType::LargeUtf8, DataType::LargeUtf8) => {
    167
    0
                    split_part_impl::<&GenericStringArray<i64>, &GenericStringArray<i64>, i64>(
    168
    0
                        args[0].as_string::<i64>(),
    169
    0
                        args[1].as_string::<i64>(),
    170
    0
                        n_array,
    171
                    )
    172
                }
    173
                (DataType::Utf8, DataType::LargeUtf8) => {
    174
    0
                    split_part_impl::<&GenericStringArray<i32>, &GenericStringArray<i64>, i32>(
    175
    0
                        args[0].as_string::<i32>(),
    176
    0
                        args[1].as_string::<i64>(),
    177
    0
                        n_array,
    178
                    )
    179
                }
    180
                (DataType::LargeUtf8, DataType::Utf8) => {
    181
    0
                    split_part_impl::<&GenericStringArray<i64>, &GenericStringArray<i32>, i64>(
    182
    0
                        args[0].as_string::<i64>(),
    183
    0
                        args[1].as_string::<i32>(),
    184
    0
                        n_array,
    185
                    )
    186
                }
    187
    0
                _ => exec_err!("Unsupported combination of argument types for split_part"),
    188
            };
    189
    0
            if is_scalar {
    190
                // If all inputs are scalar, keep the output as scalar
    191
    0
                let result = result.and_then(|arr| ScalarValue::try_from_array(&arr, 0));
    192
    0
                result.map(ColumnarValue::Scalar)
    193
            } else {
    194
    0
                result.map(ColumnarValue::Array)
    195
            }
    196
    0
        }
    197
    198
    0
        fn documentation(&self) -> Option<&Documentation> {
    199
    0
            self.doc()
    200
    0
        }
    201
    }
    202
    203
    /// impl
    204
    0
    pub fn split_part_impl<'a, StringArrType, DelimiterArrType, StringArrayLen>(
    205
    0
        string_array: StringArrType,
    206
    0
        delimiter_array: DelimiterArrType,
    207
    0
        n_array: &Int64Array,
    208
    0
    ) -> Result<ArrayRef>
    209
    0
    where
    210
    0
        StringArrType: StringArrayType<'a>,
    211
    0
        DelimiterArrType: StringArrayType<'a>,
    212
    0
        StringArrayLen: OffsetSizeTrait,
    213
    {
    214
    0
        let mut builder: GenericStringBuilder<StringArrayLen> = GenericStringBuilder::new();
    215
    216
    0
        string_array
    217
    0
            .iter()
    218
    0
            .zip(delimiter_array.iter())
    219
    0
            .zip(n_array.iter())
    220
    0
            .try_for_each(|((string, delimiter), n)| -> Result<(), DataFusionError> {
    221
    0
                match (string, delimiter, n) {
    222
    0
                    (Some(string), Some(delimiter), Some(n)) => {
    223
    0
                        let split_string: Vec<&str> = string.split(delimiter).collect();
    224
    0
                        let len = split_string.len();
    225
    226
    0
                        let index = match n.cmp(&0) {
    227
    0
                            std::cmp::Ordering::Less => len as i64 + n,
    228
                            std::cmp::Ordering::Equal => {
    229
    0
                                return exec_err!("field position must not be zero");
    230
                            }
    231
    0
                            std::cmp::Ordering::Greater => n - 1,
    232
                        } as usize;
    233
    234
    0
                        if index < len {
    235
    0
                            builder.append_value(split_string[index]);
    236
    0
                        } else {
    237
    0
                            builder.append_value("");
    238
    0
                        }
    239
                    }
    240
    0
                    _ => builder.append_null(),
    241
                }
    242
    0
                Ok(())
    243
    0
            })?;
    244
    245
    0
        Ok(Arc::new(builder.finish()) as ArrayRef)
    246
    0
    }
    247
    248
    #[cfg(test)]
    249
    mod tests {
    250
        use arrow::array::{Array, StringArray};
    251
        use arrow::datatypes::DataType::Utf8;
    252
    253
        use datafusion_common::ScalarValue;
    254
        use datafusion_common::{exec_err, Result};
    255
        use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
    256
    257
        use crate::string::split_part::SplitPartFunc;
    258
        use crate::utils::test::test_function;
    259
    260
        #[test]
    261
        fn test_functions() -> Result<()> {
    262
            test_function!(
    263
                SplitPartFunc::new(),
    264
                vec![
    265
                    ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from(
    266
                        "abc~@~def~@~ghi"
    267
                    )))),
    268
                    ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("~@~")))),
    269
                    ColumnarValue::Scalar(ScalarValue::Int64(Some(2))),
    270
                ],
    271
                Ok(Some("def")),
    272
                &str,
    273
                Utf8,
    274
                StringArray
    275
            );
    276
            test_function!(
    277
                SplitPartFunc::new(),
    278
                vec![
    279
                    ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from(
    280
                        "abc~@~def~@~ghi"
    281
                    )))),
    282
                    ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("~@~")))),
    283
                    ColumnarValue::Scalar(ScalarValue::Int64(Some(20))),
    284
                ],
    285
                Ok(Some("")),
    286
                &str,
    287
                Utf8,
    288
                StringArray
    289
            );
    290
            test_function!(
    291
                SplitPartFunc::new(),
    292
                vec![
    293
                    ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from(
    294
                        "abc~@~def~@~ghi"
    295
                    )))),
    296
                    ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("~@~")))),
    297
                    ColumnarValue::Scalar(ScalarValue::Int64(Some(-1))),
    298
                ],
    299
                Ok(Some("ghi")),
    300
                &str,
    301
                Utf8,
    302
                StringArray
    303
            );
    304
            test_function!(
    305
                SplitPartFunc::new(),
    306
                vec![
    307
                    ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from(
    308
                        "abc~@~def~@~ghi"
    309
                    )))),
    310
                    ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("~@~")))),
    311
                    ColumnarValue::Scalar(ScalarValue::Int64(Some(0))),
    312
                ],
    313
                exec_err!("field position must not be zero"),
    314
                &str,
    315
                Utf8,
    316
                StringArray
    317
            );
    318
    319
            Ok(())
    320
        }
    321
    }