## 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/unicode/substrindex.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 std::any::Any;
    19
    use std::sync::Arc;
    20
    21
    use arrow::array::{
    22
        ArrayAccessor, ArrayIter, ArrayRef, ArrowPrimitiveType, AsArray, OffsetSizeTrait,
    23
        PrimitiveArray, StringBuilder,
    24
    };
    25
    use arrow::datatypes::{DataType, Int32Type, Int64Type};
    26
    27
    use crate::utils::{make_scalar_function, utf8_to_str_type};
    28
    use datafusion_common::{exec_err, utils::take_function_args, Result};
    29
    use datafusion_expr::TypeSignature::Exact;
    30
    use datafusion_expr::{
    31
        ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility,
    32
    };
    33
    use datafusion_macros::user_doc;
    34
    35
    #[user_doc(
    36
        doc_section(label = "String Functions"),
    37
        description = r#"Returns the substring from str before count occurrences of the delimiter delim.
    38
    If count is positive, everything to the left of the final delimiter (counting from the left) is returned.
    39
    If count is negative, everything to the right of the final delimiter (counting from the right) is returned."#,
    40
        syntax_example = "substr_index(str, delim, count)",
    41
        sql_example = r#"```sql
    42
    > select substr_index('www.apache.org', '.', 1);
    43
    +---------------------------------------------------------+
    44
    | substr_index(Utf8("www.apache.org"),Utf8("."),Int64(1)) |
    45
    +---------------------------------------------------------+
    46
    | www                                                     |
    47
    +---------------------------------------------------------+
    48
    > select substr_index('www.apache.org', '.', -1);
    49
    +----------------------------------------------------------+
    50
    | substr_index(Utf8("www.apache.org"),Utf8("."),Int64(-1)) |
    51
    +----------------------------------------------------------+
    52
    | org                                                      |
    53
    +----------------------------------------------------------+
    54
    ```"#,
    55
        standard_argument(name = "str", prefix = "String"),
    56
        argument(
    57
            name = "delim",
    58
            description = "The string to find in str to split str."
    59
        ),
    60
        argument(
    61
            name = "count",
    62
            description = "The number of times to search for the delimiter. Can be either a positive or negative number."
    63
        )
    64
    )]
    65
    #[derive(Debug, PartialEq, Eq, Hash)]
    66
    pub struct SubstrIndexFunc {
    67
        signature: Signature,
    68
        aliases: Vec<String>,
    69
    }
    70
    71
    impl Default for SubstrIndexFunc {
    72
    0
        fn default() -> Self {
    73
    0
            Self::new()
    74
    0
        }
    75
    }
    76
    77
    impl SubstrIndexFunc {
    78
    0
        pub fn new() -> Self {
    79
            use DataType::*;
    80
    0
            Self {
    81
    0
                signature: Signature::one_of(
    82
    0
                    vec![
    83
    0
                        Exact(vec![Utf8View, Utf8View, Int64]),
    84
    0
                        Exact(vec![Utf8, Utf8, Int64]),
    85
    0
                        Exact(vec![LargeUtf8, LargeUtf8, Int64]),
    86
    0
                    ],
    87
    0
                    Volatility::Immutable,
    88
    0
                ),
    89
    0
                aliases: vec![String::from("substring_index")],
    90
    0
            }
    91
    0
        }
    92
    }
    93
    94
    impl ScalarUDFImpl for SubstrIndexFunc {
    95
    0
        fn as_any(&self) -> &dyn Any {
    96
    0
            self
    97
    0
        }
    98
    99
    0
        fn name(&self) -> &str {
    100
    0
            "substr_index"
    101
    0
        }
    102
    103
    0
        fn signature(&self) -> &Signature {
    104
    0
            &self.signature
    105
    0
        }
    106
    107
    0
        fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
    108
    0
            utf8_to_str_type(&arg_types[0], "substr_index")
    109
    0
        }
    110
    111
    0
        fn invoke_with_args(
    112
    0
            &self,
    113
    0
            args: datafusion_expr::ScalarFunctionArgs,
    114
    0
        ) -> Result<ColumnarValue> {
    115
    0
            make_scalar_function(substr_index, vec![])(&args.args)
    116
    0
        }
    117
    118
    0
        fn aliases(&self) -> &[String] {
    119
    0
            &self.aliases
    120
    0
        }
    121
    122
    0
        fn documentation(&self) -> Option<&Documentation> {
    123
    0
            self.doc()
    124
    0
        }
    125
    }
    126
    127
    /// Returns the substring from str before count occurrences of the delimiter delim. If count is positive, everything to the left of the final delimiter (counting from the left) is returned. If count is negative, everything to the right of the final delimiter (counting from the right) is returned.
    128
    /// SUBSTRING_INDEX('www.apache.org', '.', 1) = www
    129
    /// SUBSTRING_INDEX('www.apache.org', '.', 2) = www.apache
    130
    /// SUBSTRING_INDEX('www.apache.org', '.', -2) = apache.org
    131
    /// SUBSTRING_INDEX('www.apache.org', '.', -1) = org
    132
    0
    fn substr_index(args: &[ArrayRef]) -> Result<ArrayRef> {
    133
    0
        let [str, delim, count] = take_function_args("substr_index", args)?;
    134
    135
    0
        match str.data_type() {
    136
            DataType::Utf8 => {
    137
    0
                let string_array = str.as_string::<i32>();
    138
    0
                let delimiter_array = delim.as_string::<i32>();
    139
    0
                let count_array: &PrimitiveArray<Int64Type> = count.as_primitive();
    140
    0
                substr_index_general::<Int32Type, _, _>(
    141
    0
                    string_array,
    142
    0
                    delimiter_array,
    143
    0
                    count_array,
    144
                )
    145
            }
    146
            DataType::LargeUtf8 => {
    147
    0
                let string_array = str.as_string::<i64>();
    148
    0
                let delimiter_array = delim.as_string::<i64>();
    149
    0
                let count_array: &PrimitiveArray<Int64Type> = count.as_primitive();
    150
    0
                substr_index_general::<Int64Type, _, _>(
    151
    0
                    string_array,
    152
    0
                    delimiter_array,
    153
    0
                    count_array,
    154
                )
    155
            }
    156
            DataType::Utf8View => {
    157
    0
                let string_array = str.as_string_view();
    158
    0
                let delimiter_array = delim.as_string_view();
    159
    0
                let count_array: &PrimitiveArray<Int64Type> = count.as_primitive();
    160
    0
                substr_index_general::<Int32Type, _, _>(
    161
    0
                    string_array,
    162
    0
                    delimiter_array,
    163
    0
                    count_array,
    164
                )
    165
            }
    166
    0
            other => {
    167
    0
                exec_err!("Unsupported data type {other:?} for function substr_index")
    168
            }
    169
        }
    170
    0
    }
    171
    172
    0
    pub fn substr_index_general<
    173
    0
        'a,
    174
    0
        T: ArrowPrimitiveType,
    175
    0
        V: ArrayAccessor<Item = &'a str>,
    176
    0
        P: ArrayAccessor<Item = i64>,
    177
    0
    >(
    178
    0
        string_array: V,
    179
    0
        delimiter_array: V,
    180
    0
        count_array: P,
    181
    0
    ) -> Result<ArrayRef>
    182
    0
    where
    183
    0
        T::Native: OffsetSizeTrait,
    184
    {
    185
    0
        let mut builder = StringBuilder::new();
    186
    0
        let string_iter = ArrayIter::new(string_array);
    187
    0
        let delimiter_array_iter = ArrayIter::new(delimiter_array);
    188
    0
        let count_array_iter = ArrayIter::new(count_array);
    189
    0
        string_iter
    190
    0
            .zip(delimiter_array_iter)
    191
    0
            .zip(count_array_iter)
    192
    0
            .for_each(|((string, delimiter), n)| match (string, delimiter, n) {
    193
    0
                (Some(string), Some(delimiter), Some(n)) => {
    194
                    // In MySQL, these cases will return an empty string.
    195
    0
                    if n == 0 || string.is_empty() || delimiter.is_empty() {
    196
    0
                        builder.append_value("");
    197
    0
                        return;
    198
    0
                    }
    199
    200
    0
                    let occurrences = usize::try_from(n.unsigned_abs()).unwrap_or(usize::MAX);
    201
    0
                    let length = if n > 0 {
    202
    0
                        let split = string.split(delimiter);
    203
    0
                        split
    204
    0
                            .take(occurrences)
    205
    0
                            .map(|s| s.len() + delimiter.len())
    206
    0
                            .sum::<usize>()
    207
    0
                            - delimiter.len()
    208
                    } else {
    209
    0
                        let split = string.rsplit(delimiter);
    210
    0
                        split
    211
    0
                            .take(occurrences)
    212
    0
                            .map(|s| s.len() + delimiter.len())
    213
    0
                            .sum::<usize>()
    214
    0
                            - delimiter.len()
    215
                    };
    216
    0
                    if n > 0 {
    217
    0
                        match string.get(..length) {
    218
    0
                            Some(substring) => builder.append_value(substring),
    219
    0
                            None => builder.append_null(),
    220
                        }
    221
                    } else {
    222
    0
                        match string.get(string.len().saturating_sub(length)..) {
    223
    0
                            Some(substring) => builder.append_value(substring),
    224
    0
                            None => builder.append_null(),
    225
                        }
    226
                    }
    227
                }
    228
    0
                _ => builder.append_null(),
    229
    0
            });
    230
    231
    0
        Ok(Arc::new(builder.finish()) as ArrayRef)
    232
    0
    }
    233
    234
    #[cfg(test)]
    235
    mod tests {
    236
        use arrow::array::{Array, StringArray};
    237
        use arrow::datatypes::DataType::Utf8;
    238
    239
        use datafusion_common::{Result, ScalarValue};
    240
        use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
    241
    242
        use crate::unicode::substrindex::SubstrIndexFunc;
    243
        use crate::utils::test::test_function;
    244
    245
        #[test]
    246
        fn test_functions() -> Result<()> {
    247
            test_function!(
    248
                SubstrIndexFunc::new(),
    249
                vec![
    250
                    ColumnarValue::Scalar(ScalarValue::from("www.apache.org")),
    251
                    ColumnarValue::Scalar(ScalarValue::from(".")),
    252
                    ColumnarValue::Scalar(ScalarValue::from(1i64)),
    253
                ],
    254
                Ok(Some("www")),
    255
                &str,
    256
                Utf8,
    257
                StringArray
    258
            );
    259
            test_function!(
    260
                SubstrIndexFunc::new(),
    261
                vec![
    262
                    ColumnarValue::Scalar(ScalarValue::from("www.apache.org")),
    263
                    ColumnarValue::Scalar(ScalarValue::from(".")),
    264
                    ColumnarValue::Scalar(ScalarValue::from(2i64)),
    265
                ],
    266
                Ok(Some("www.apache")),
    267
                &str,
    268
                Utf8,
    269
                StringArray
    270
            );
    271
            test_function!(
    272
                SubstrIndexFunc::new(),
    273
                vec![
    274
                    ColumnarValue::Scalar(ScalarValue::from("www.apache.org")),
    275
                    ColumnarValue::Scalar(ScalarValue::from(".")),
    276
                    ColumnarValue::Scalar(ScalarValue::from(-2i64)),
    277
                ],
    278
                Ok(Some("apache.org")),
    279
                &str,
    280
                Utf8,
    281
                StringArray
    282
            );
    283
            test_function!(
    284
                SubstrIndexFunc::new(),
    285
                vec![
    286
                    ColumnarValue::Scalar(ScalarValue::from("www.apache.org")),
    287
                    ColumnarValue::Scalar(ScalarValue::from(".")),
    288
                    ColumnarValue::Scalar(ScalarValue::from(-1i64)),
    289
                ],
    290
                Ok(Some("org")),
    291
                &str,
    292
                Utf8,
    293
                StringArray
    294
            );
    295
            test_function!(
    296
                SubstrIndexFunc::new(),
    297
                vec![
    298
                    ColumnarValue::Scalar(ScalarValue::from("www.apache.org")),
    299
                    ColumnarValue::Scalar(ScalarValue::from(".")),
    300
                    ColumnarValue::Scalar(ScalarValue::from(0i64)),
    301
                ],
    302
                Ok(Some("")),
    303
                &str,
    304
                Utf8,
    305
                StringArray
    306
            );
    307
            test_function!(
    308
                SubstrIndexFunc::new(),
    309
                vec![
    310
                    ColumnarValue::Scalar(ScalarValue::from("")),
    311
                    ColumnarValue::Scalar(ScalarValue::from(".")),
    312
                    ColumnarValue::Scalar(ScalarValue::from(1i64)),
    313
                ],
    314
                Ok(Some("")),
    315
                &str,
    316
                Utf8,
    317
                StringArray
    318
            );
    319
            test_function!(
    320
                SubstrIndexFunc::new(),
    321
                vec![
    322
                    ColumnarValue::Scalar(ScalarValue::from("www.apache.org")),
    323
                    ColumnarValue::Scalar(ScalarValue::from("")),
    324
                    ColumnarValue::Scalar(ScalarValue::from(1i64)),
    325
                ],
    326
                Ok(Some("")),
    327
                &str,
    328
                Utf8,
    329
                StringArray
    330
            );
    331
    332
            Ok(())
    333
        }
    334
    }