## 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/regex/regexpcount.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::regex::{compile_and_cache_regex, compile_regex};
    19
    use arrow::array::{Array, ArrayRef, AsArray, Datum, Int64Array, StringArrayType};
    20
    use arrow::datatypes::{DataType, Int64Type};
    21
    use arrow::datatypes::{
    22
        DataType::Int64, DataType::LargeUtf8, DataType::Utf8, DataType::Utf8View,
    23
    };
    24
    use arrow::error::ArrowError;
    25
    use datafusion_common::{exec_err, internal_err, Result, ScalarValue};
    26
    use datafusion_expr::{
    27
        ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignature::Exact,
    28
        TypeSignature::Uniform, Volatility,
    29
    };
    30
    use datafusion_macros::user_doc;
    31
    use itertools::izip;
    32
    use regex::Regex;
    33
    use std::collections::HashMap;
    34
    use std::sync::Arc;
    35
    36
    #[user_doc(
    37
        doc_section(label = "Regular Expression Functions"),
    38
        description = "Returns the number of matches that a [regular expression](https://docs.rs/regex/latest/regex/#syntax) has in a string.",
    39
        syntax_example = "regexp_count(str, regexp[, start, flags])",
    40
        sql_example = r#"```sql
    41
    > select regexp_count('abcAbAbc', 'abc', 2, 'i');
    42
    +---------------------------------------------------------------+
    43
    | regexp_count(Utf8("abcAbAbc"),Utf8("abc"),Int64(2),Utf8("i")) |
    44
    +---------------------------------------------------------------+
    45
    | 1                                                             |
    46
    +---------------------------------------------------------------+
    47
    ```"#,
    48
        standard_argument(name = "str", prefix = "String"),
    49
        standard_argument(name = "regexp", prefix = "Regular"),
    50
        argument(
    51
            name = "start",
    52
            description = "- **start**: Optional start position (the first position is 1) to search for the regular expression. Can be a constant, column, or function."
    53
        ),
    54
        argument(
    55
            name = "flags",
    56
            description = r#"Optional regular expression flags that control the behavior of the regular expression. The following flags are supported:
    57
      - **i**: case-insensitive: letters match both upper and lower case
    58
      - **m**: multi-line mode: ^ and $ match begin/end of line
    59
      - **s**: allow . to match \n
    60
      - **R**: enables CRLF mode: when multi-line mode is enabled, \r\n is used
    61
      - **U**: swap the meaning of x* and x*?"#
    62
        )
    63
    )]
    64
    #[derive(Debug, PartialEq, Eq, Hash)]
    65
    pub struct RegexpCountFunc {
    66
        signature: Signature,
    67
    }
    68
    69
    impl Default for RegexpCountFunc {
    70
    0
        fn default() -> Self {
    71
    0
            Self::new()
    72
    0
        }
    73
    }
    74
    75
    impl RegexpCountFunc {
    76
    0
        pub fn new() -> Self {
    77
    0
            Self {
    78
    0
                signature: Signature::one_of(
    79
    0
                    vec![
    80
    0
                        Uniform(2, vec![Utf8View, LargeUtf8, Utf8]),
    81
    0
                        Exact(vec![Utf8View, Utf8View, Int64]),
    82
    0
                        Exact(vec![LargeUtf8, LargeUtf8, Int64]),
    83
    0
                        Exact(vec![Utf8, Utf8, Int64]),
    84
    0
                        Exact(vec![Utf8View, Utf8View, Int64, Utf8View]),
    85
    0
                        Exact(vec![LargeUtf8, LargeUtf8, Int64, LargeUtf8]),
    86
    0
                        Exact(vec![Utf8, Utf8, Int64, Utf8]),
    87
    0
                    ],
    88
    0
                    Volatility::Immutable,
    89
    0
                ),
    90
    0
            }
    91
    0
        }
    92
    }
    93
    94
    impl ScalarUDFImpl for RegexpCountFunc {
    95
    0
        fn as_any(&self) -> &dyn std::any::Any {
    96
    0
            self
    97
    0
        }
    98
    99
    0
        fn name(&self) -> &str {
    100
    0
            "regexp_count"
    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
            Ok(Int64)
    109
    0
        }
    110
    111
    0
        fn invoke_with_args(
    112
    0
            &self,
    113
    0
            args: datafusion_expr::ScalarFunctionArgs,
    114
    0
        ) -> Result<ColumnarValue> {
    115
    0
            let args = &args.args;
    116
    117
    0
            let len = args
    118
    0
                .iter()
    119
    0
                .fold(Option::<usize>::None, |acc, arg| match arg {
    120
    0
                    ColumnarValue::Scalar(_) => acc,
    121
    0
                    ColumnarValue::Array(a) => Some(a.len()),
    122
    0
                });
    123
    124
    0
            let is_scalar = len.is_none();
    125
    0
            let inferred_length = len.unwrap_or(1);
    126
    0
            let args = args
    127
    0
                .iter()
    128
    0
                .map(|arg| arg.to_array(inferred_length))
    129
    0
                .collect::<Result<Vec<_>>>()?;
    130
    131
    0
            let result = regexp_count_func(&args);
    132
    0
            if is_scalar {
    133
                // If all inputs are scalar, keeps output as scalar
    134
    0
                let result = result.and_then(|arr| ScalarValue::try_from_array(&arr, 0));
    135
    0
                result.map(ColumnarValue::Scalar)
    136
            } else {
    137
    0
                result.map(ColumnarValue::Array)
    138
            }
    139
    0
        }
    140
    141
    0
        fn documentation(&self) -> Option<&Documentation> {
    142
    0
            self.doc()
    143
    0
        }
    144
    }
    145
    146
    0
    pub fn regexp_count_func(args: &[ArrayRef]) -> Result<ArrayRef> {
    147
    0
        let args_len = args.len();
    148
    0
        if !(2..=4).contains(&args_len) {
    149
    0
            return exec_err!("regexp_count was called with {args_len} arguments. It requires at least 2 and at most 4.");
    150
    0
        }
    151
    152
    0
        let values = &args[0];
    153
    0
        match values.data_type() {
    154
    0
            Utf8 | LargeUtf8 | Utf8View => (),
    155
    0
            other => {
    156
    0
                return internal_err!(
    157
    0
                    "Unsupported data type {other:?} for function regexp_count"
    158
                );
    159
            }
    160
        }
    161
    162
    0
        regexp_count(
    163
    0
            values,
    164
    0
            &args[1],
    165
    0
            if args_len > 2 { Some(&args[2]) } else { None },
    166
    0
            if args_len > 3 { Some(&args[3]) } else { None },
    167
        )
    168
    0
        .map_err(|e| e.into())
    169
    0
    }
    170
    171
    /// `arrow-rs` style implementation of `regexp_count` function.
    172
    /// This function `regexp_count` is responsible for counting the occurrences of a regular expression pattern
    173
    /// within a string array. It supports optional start positions and flags for case insensitivity.
    174
    ///
    175
    /// The function accepts a variable number of arguments:
    176
    /// - `values`: The array of strings to search within.
    177
    /// - `regex_array`: The array of regular expression patterns to search for.
    178
    /// - `start_array` (optional): The array of start positions for the search.
    179
    /// - `flags_array` (optional): The array of flags to modify the search behavior (e.g., case insensitivity).
    180
    ///
    181
    /// The function handles different combinations of scalar and array inputs for the regex patterns, start positions,
    182
    /// and flags. It uses a cache to store compiled regular expressions for efficiency.
    183
    ///
    184
    /// # Errors
    185
    /// Returns an error if the input arrays have mismatched lengths or if the regular expression fails to compile.
    186
    0
    pub fn regexp_count(
    187
    0
        values: &dyn Array,
    188
    0
        regex_array: &dyn Datum,
    189
    0
        start_array: Option<&dyn Datum>,
    190
    0
        flags_array: Option<&dyn Datum>,
    191
    0
    ) -> Result<ArrayRef, ArrowError> {
    192
    0
        let (regex_array, is_regex_scalar) = regex_array.get();
    193
    0
        let (start_array, is_start_scalar) = start_array.map_or((None, true), |start| {
    194
    0
            let (start, is_start_scalar) = start.get();
    195
    0
            (Some(start), is_start_scalar)
    196
    0
        });
    197
    0
        let (flags_array, is_flags_scalar) = flags_array.map_or((None, true), |flags| {
    198
    0
            let (flags, is_flags_scalar) = flags.get();
    199
    0
            (Some(flags), is_flags_scalar)
    200
    0
        });
    201
    202
    0
        match (values.data_type(), regex_array.data_type(), flags_array) {
    203
    0
            (Utf8, Utf8, None) => regexp_count_inner(
    204
    0
                values.as_string::<i32>(),
    205
    0
                regex_array.as_string::<i32>(),
    206
    0
                is_regex_scalar,
    207
    0
                start_array.map(|start| start.as_primitive::<Int64Type>()),
    208
    0
                is_start_scalar,
    209
    0
                None,
    210
    0
                is_flags_scalar,
    211
            ),
    212
    0
            (Utf8, Utf8, Some(flags_array)) if *flags_array.data_type() == Utf8 => regexp_count_inner(
    213
    0
                values.as_string::<i32>(),
    214
    0
                regex_array.as_string::<i32>(),
    215
    0
                is_regex_scalar,
    216
    0
                start_array.map(|start| start.as_primitive::<Int64Type>()),
    217
    0
                is_start_scalar,
    218
    0
                Some(flags_array.as_string::<i32>()),
    219
    0
                is_flags_scalar,
    220
            ),
    221
    0
            (LargeUtf8, LargeUtf8, None) => regexp_count_inner(
    222
    0
                values.as_string::<i64>(),
    223
    0
                regex_array.as_string::<i64>(),
    224
    0
                is_regex_scalar,
    225
    0
                start_array.map(|start| start.as_primitive::<Int64Type>()),
    226
    0
                is_start_scalar,
    227
    0
                None,
    228
    0
                is_flags_scalar,
    229
            ),
    230
    0
            (LargeUtf8, LargeUtf8, Some(flags_array)) if *flags_array.data_type() == LargeUtf8 => regexp_count_inner(
    231
    0
                values.as_string::<i64>(),
    232
    0
                regex_array.as_string::<i64>(),
    233
    0
                is_regex_scalar,
    234
    0
                start_array.map(|start| start.as_primitive::<Int64Type>()),
    235
    0
                is_start_scalar,
    236
    0
                Some(flags_array.as_string::<i64>()),
    237
    0
                is_flags_scalar,
    238
            ),
    239
    0
            (Utf8View, Utf8View, None) => regexp_count_inner(
    240
    0
                values.as_string_view(),
    241
    0
                regex_array.as_string_view(),
    242
    0
                is_regex_scalar,
    243
    0
                start_array.map(|start| start.as_primitive::<Int64Type>()),
    244
    0
                is_start_scalar,
    245
    0
                None,
    246
    0
                is_flags_scalar,
    247
            ),
    248
    0
            (Utf8View, Utf8View, Some(flags_array)) if *flags_array.data_type() == Utf8View => regexp_count_inner(
    249
    0
                values.as_string_view(),
    250
    0
                regex_array.as_string_view(),
    251
    0
                is_regex_scalar,
    252
    0
                start_array.map(|start| start.as_primitive::<Int64Type>()),
    253
    0
                is_start_scalar,
    254
    0
                Some(flags_array.as_string_view()),
    255
    0
                is_flags_scalar,
    256
            ),
    257
    0
            _ => Err(ArrowError::ComputeError(
    258
    0
                "regexp_count() expected the input arrays to be of type Utf8, LargeUtf8, or Utf8View and the data types of the values, regex_array, and flags_array to match".to_string(),
    259
    0
            )),
    260
        }
    261
    0
    }
    262
    263
    0
    pub fn regexp_count_inner<'a, S>(
    264
    0
        values: S,
    265
    0
        regex_array: S,
    266
    0
        is_regex_scalar: bool,
    267
    0
        start_array: Option<&Int64Array>,
    268
    0
        is_start_scalar: bool,
    269
    0
        flags_array: Option<S>,
    270
    0
        is_flags_scalar: bool,
    271
    0
    ) -> Result<ArrayRef, ArrowError>
    272
    0
    where
    273
    0
        S: StringArrayType<'a>,
    274
    {
    275
    0
        let (regex_scalar, is_regex_scalar) = if is_regex_scalar || regex_array.len() == 1 {
    276
    0
            (Some(regex_array.value(0)), true)
    277
        } else {
    278
    0
            (None, false)
    279
        };
    280
    281
    0
        let (start_array, start_scalar, is_start_scalar) =
    282
    0
            if let Some(start_array) = start_array {
    283
    0
                if is_start_scalar || start_array.len() == 1 {
    284
    0
                    (None, Some(start_array.value(0)), true)
    285
                } else {
    286
    0
                    (Some(start_array), None, false)
    287
                }
    288
            } else {
    289
    0
                (None, Some(1), true)
    290
            };
    291
    292
    0
        let (flags_array, flags_scalar, is_flags_scalar) =
    293
    0
            if let Some(flags_array) = flags_array {
    294
    0
                if is_flags_scalar || flags_array.len() == 1 {
    295
    0
                    (None, Some(flags_array.value(0)), true)
    296
                } else {
    297
    0
                    (Some(flags_array), None, false)
    298
                }
    299
            } else {
    300
    0
                (None, None, true)
    301
            };
    302
    303
    0
        let mut regex_cache = HashMap::new();
    304
    305
    0
        match (is_regex_scalar, is_start_scalar, is_flags_scalar) {
    306
            (true, true, true) => {
    307
    0
                let regex = match regex_scalar {
    308
    0
                    None | Some("") => {
    309
    0
                        return Ok(Arc::new(Int64Array::from(vec![0; values.len()])))
    310
                    }
    311
    0
                    Some(regex) => regex,
    312
                };
    313
    314
    0
                let pattern = compile_regex(regex, flags_scalar)?;
    315
    316
    0
                Ok(Arc::new(
    317
    0
                    values
    318
    0
                        .iter()
    319
    0
                        .map(|value| count_matches(value, &pattern, start_scalar))
    320
    0
                        .collect::<Result<Int64Array, ArrowError>>()?,
    321
                ))
    322
            }
    323
            (true, true, false) => {
    324
    0
                let regex = match regex_scalar {
    325
    0
                    None | Some("") => {
    326
    0
                        return Ok(Arc::new(Int64Array::from(vec![0; values.len()])))
    327
                    }
    328
    0
                    Some(regex) => regex,
    329
                };
    330
    331
    0
                let flags_array = flags_array.unwrap();
    332
    0
                if values.len() != flags_array.len() {
    333
    0
                    return Err(ArrowError::ComputeError(format!(
    334
    0
                        "flags_array must be the same length as values array; got {} and {}",
    335
    0
                        flags_array.len(),
    336
    0
                        values.len(),
    337
    0
                    )));
    338
    0
                }
    339
    340
    0
                Ok(Arc::new(
    341
    0
                    values
    342
    0
                        .iter()
    343
    0
                        .zip(flags_array.iter())
    344
    0
                        .map(|(value, flags)| {
    345
    0
                            let pattern =
    346
    0
                                compile_and_cache_regex(regex, flags, &mut regex_cache)?;
    347
    0
                            count_matches(value, pattern, start_scalar)
    348
    0
                        })
    349
    0
                        .collect::<Result<Int64Array, ArrowError>>()?,
    350
                ))
    351
            }
    352
            (true, false, true) => {
    353
    0
                let regex = match regex_scalar {
    354
    0
                    None | Some("") => {
    355
    0
                        return Ok(Arc::new(Int64Array::from(vec![0; values.len()])))
    356
                    }
    357
    0
                    Some(regex) => regex,
    358
                };
    359
    360
    0
                let pattern = compile_regex(regex, flags_scalar)?;
    361
    362
    0
                let start_array = start_array.unwrap();
    363
    364
    0
                Ok(Arc::new(
    365
    0
                    values
    366
    0
                        .iter()
    367
    0
                        .zip(start_array.iter())
    368
    0
                        .map(|(value, start)| count_matches(value, &pattern, start))
    369
    0
                        .collect::<Result<Int64Array, ArrowError>>()?,
    370
                ))
    371
            }
    372
            (true, false, false) => {
    373
    0
                let regex = match regex_scalar {
    374
    0
                    None | Some("") => {
    375
    0
                        return Ok(Arc::new(Int64Array::from(vec![0; values.len()])))
    376
                    }
    377
    0
                    Some(regex) => regex,
    378
                };
    379
    380
    0
                let flags_array = flags_array.unwrap();
    381
    0
                if values.len() != flags_array.len() {
    382
    0
                    return Err(ArrowError::ComputeError(format!(
    383
    0
                        "flags_array must be the same length as values array; got {} and {}",
    384
    0
                        flags_array.len(),
    385
    0
                        values.len(),
    386
    0
                    )));
    387
    0
                }
    388
    389
    0
                Ok(Arc::new(
    390
    0
                    izip!(
    391
    0
                        values.iter(),
    392
    0
                        start_array.unwrap().iter(),
    393
    0
                        flags_array.iter()
    394
                    )
    395
    0
                    .map(|(value, start, flags)| {
    396
    0
                        let pattern =
    397
    0
                            compile_and_cache_regex(regex, flags, &mut regex_cache)?;
    398
    399
    0
                        count_matches(value, pattern, start)
    400
    0
                    })
    401
    0
                    .collect::<Result<Int64Array, ArrowError>>()?,
    402
                ))
    403
            }
    404
            (false, true, true) => {
    405
    0
                if values.len() != regex_array.len() {
    406
    0
                    return Err(ArrowError::ComputeError(format!(
    407
    0
                        "regex_array must be the same length as values array; got {} and {}",
    408
    0
                        regex_array.len(),
    409
    0
                        values.len(),
    410
    0
                    )));
    411
    0
                }
    412
    413
    0
                Ok(Arc::new(
    414
    0
                    values
    415
    0
                        .iter()
    416
    0
                        .zip(regex_array.iter())
    417
    0
                        .map(|(value, regex)| {
    418
    0
                            let regex = match regex {
    419
    0
                                None | Some("") => return Ok(0),
    420
    0
                                Some(regex) => regex,
    421
                            };
    422
    423
    0
                            let pattern = compile_and_cache_regex(
    424
    0
                                regex,
    425
    0
                                flags_scalar,
    426
    0
                                &mut regex_cache,
    427
    0
                            )?;
    428
    0
                            count_matches(value, pattern, start_scalar)
    429
    0
                        })
    430
    0
                        .collect::<Result<Int64Array, ArrowError>>()?,
    431
                ))
    432
            }
    433
            (false, true, false) => {
    434
    0
                if values.len() != regex_array.len() {
    435
    0
                    return Err(ArrowError::ComputeError(format!(
    436
    0
                        "regex_array must be the same length as values array; got {} and {}",
    437
    0
                        regex_array.len(),
    438
    0
                        values.len(),
    439
    0
                    )));
    440
    0
                }
    441
    442
    0
                let flags_array = flags_array.unwrap();
    443
    0
                if values.len() != flags_array.len() {
    444
    0
                    return Err(ArrowError::ComputeError(format!(
    445
    0
                        "flags_array must be the same length as values array; got {} and {}",
    446
    0
                        flags_array.len(),
    447
    0
                        values.len(),
    448
    0
                    )));
    449
    0
                }
    450
    451
    0
                Ok(Arc::new(
    452
    0
                    izip!(values.iter(), regex_array.iter(), flags_array.iter())
    453
    0
                        .map(|(value, regex, flags)| {
    454
    0
                            let regex = match regex {
    455
    0
                                None | Some("") => return Ok(0),
    456
    0
                                Some(regex) => regex,
    457
                            };
    458
    459
    0
                            let pattern =
    460
    0
                                compile_and_cache_regex(regex, flags, &mut regex_cache)?;
    461
    462
    0
                            count_matches(value, pattern, start_scalar)
    463
    0
                        })
    464
    0
                        .collect::<Result<Int64Array, ArrowError>>()?,
    465
                ))
    466
            }
    467
            (false, false, true) => {
    468
    0
                if values.len() != regex_array.len() {
    469
    0
                    return Err(ArrowError::ComputeError(format!(
    470
    0
                        "regex_array must be the same length as values array; got {} and {}",
    471
    0
                        regex_array.len(),
    472
    0
                        values.len(),
    473
    0
                    )));
    474
    0
                }
    475
    476
    0
                let start_array = start_array.unwrap();
    477
    0
                if values.len() != start_array.len() {
    478
    0
                    return Err(ArrowError::ComputeError(format!(
    479
    0
                        "start_array must be the same length as values array; got {} and {}",
    480
    0
                        start_array.len(),
    481
    0
                        values.len(),
    482
    0
                    )));
    483
    0
                }
    484
    485
    0
                Ok(Arc::new(
    486
    0
                    izip!(values.iter(), regex_array.iter(), start_array.iter())
    487
    0
                        .map(|(value, regex, start)| {
    488
    0
                            let regex = match regex {
    489
    0
                                None | Some("") => return Ok(0),
    490
    0
                                Some(regex) => regex,
    491
                            };
    492
    493
    0
                            let pattern = compile_and_cache_regex(
    494
    0
                                regex,
    495
    0
                                flags_scalar,
    496
    0
                                &mut regex_cache,
    497
    0
                            )?;
    498
    0
                            count_matches(value, pattern, start)
    499
    0
                        })
    500
    0
                        .collect::<Result<Int64Array, ArrowError>>()?,
    501
                ))
    502
            }
    503
            (false, false, false) => {
    504
    0
                if values.len() != regex_array.len() {
    505
    0
                    return Err(ArrowError::ComputeError(format!(
    506
    0
                        "regex_array must be the same length as values array; got {} and {}",
    507
    0
                        regex_array.len(),
    508
    0
                        values.len(),
    509
    0
                    )));
    510
    0
                }
    511
    512
    0
                let start_array = start_array.unwrap();
    513
    0
                if values.len() != start_array.len() {
    514
    0
                    return Err(ArrowError::ComputeError(format!(
    515
    0
                        "start_array must be the same length as values array; got {} and {}",
    516
    0
                        start_array.len(),
    517
    0
                        values.len(),
    518
    0
                    )));
    519
    0
                }
    520
    521
    0
                let flags_array = flags_array.unwrap();
    522
    0
                if values.len() != flags_array.len() {
    523
    0
                    return Err(ArrowError::ComputeError(format!(
    524
    0
                        "flags_array must be the same length as values array; got {} and {}",
    525
    0
                        flags_array.len(),
    526
    0
                        values.len(),
    527
    0
                    )));
    528
    0
                }
    529
    530
    0
                Ok(Arc::new(
    531
    0
                    izip!(
    532
    0
                        values.iter(),
    533
    0
                        regex_array.iter(),
    534
    0
                        start_array.iter(),
    535
    0
                        flags_array.iter()
    536
                    )
    537
    0
                    .map(|(value, regex, start, flags)| {
    538
    0
                        let regex = match regex {
    539
    0
                            None | Some("") => return Ok(0),
    540
    0
                            Some(regex) => regex,
    541
                        };
    542
    543
    0
                        let pattern =
    544
    0
                            compile_and_cache_regex(regex, flags, &mut regex_cache)?;
    545
    0
                        count_matches(value, pattern, start)
    546
    0
                    })
    547
    0
                    .collect::<Result<Int64Array, ArrowError>>()?,
    548
                ))
    549
            }
    550
        }
    551
    0
    }
    552
    553
    0
    fn count_matches(
    554
    0
        value: Option<&str>,
    555
    0
        pattern: &Regex,
    556
    0
        start: Option<i64>,
    557
    0
    ) -> Result<i64, ArrowError> {
    558
    0
        let value = match value {
    559
    0
            None | Some("") => return Ok(0),
    560
    0
            Some(value) => value,
    561
        };
    562
    563
    0
        if let Some(start) = start {
    564
    0
            if start < 1 {
    565
    0
                return Err(ArrowError::ComputeError(
    566
    0
                    "regexp_count() requires start to be 1 based".to_string(),
    567
    0
                ));
    568
    0
            }
    569
    570
    0
            let find_slice = value.chars().skip(start as usize - 1).collect::<String>();
    571
    0
            let count = pattern.find_iter(find_slice.as_str()).count();
    572
    0
            Ok(count as i64)
    573
        } else {
    574
    0
            let count = pattern.find_iter(value).count();
    575
    0
            Ok(count as i64)
    576
        }
    577
    0
    }
    578
    579
    #[cfg(test)]
    580
    mod tests {
    581
        use super::*;
    582
        use arrow::array::{GenericStringArray, StringViewArray};
    583
        use arrow::datatypes::Field;
    584
        use datafusion_common::config::ConfigOptions;
    585
        use datafusion_expr::ScalarFunctionArgs;
    586
    587
        #[test]
    588
        fn test_regexp_count() {
    589
            test_case_sensitive_regexp_count_scalar();
    590
            test_case_sensitive_regexp_count_scalar_start();
    591
            test_case_insensitive_regexp_count_scalar_flags();
    592
            test_case_sensitive_regexp_count_start_scalar_complex();
    593
    594
            test_case_sensitive_regexp_count_array::<GenericStringArray<i32>>();
    595
            test_case_sensitive_regexp_count_array::<GenericStringArray<i64>>();
    596
            test_case_sensitive_regexp_count_array::<StringViewArray>();
    597
    598
            test_case_sensitive_regexp_count_array_start::<GenericStringArray<i32>>();
    599
            test_case_sensitive_regexp_count_array_start::<GenericStringArray<i64>>();
    600
            test_case_sensitive_regexp_count_array_start::<StringViewArray>();
    601
    602
            test_case_insensitive_regexp_count_array_flags::<GenericStringArray<i32>>();
    603
            test_case_insensitive_regexp_count_array_flags::<GenericStringArray<i64>>();
    604
            test_case_insensitive_regexp_count_array_flags::<StringViewArray>();
    605
    606
            test_case_sensitive_regexp_count_array_complex::<GenericStringArray<i32>>();
    607
            test_case_sensitive_regexp_count_array_complex::<GenericStringArray<i64>>();
    608
            test_case_sensitive_regexp_count_array_complex::<StringViewArray>();
    609
    610
            test_case_regexp_count_cache_check::<GenericStringArray<i32>>();
    611
        }
    612
    613
        fn regexp_count_with_scalar_values(args: &[ScalarValue]) -> Result<ColumnarValue> {
    614
            let args_values = args
    615
                .iter()
    616
                .map(|sv| ColumnarValue::Scalar(sv.clone()))
    617
                .collect();
    618
    619
            let arg_fields = args
    620
                .iter()
    621
                .enumerate()
    622
                .map(|(idx, a)| Field::new(format!("arg_{idx}"), a.data_type(), true).into())
    623
                .collect::<Vec<_>>();
    624
    625
            RegexpCountFunc::new().invoke_with_args(ScalarFunctionArgs {
    626
                args: args_values,
    627
                arg_fields,
    628
                number_rows: args.len(),
    629
                return_field: Field::new("f", Int64, true).into(),
    630
                config_options: Arc::new(ConfigOptions::default()),
    631
            })
    632
        }
    633
    634
        fn test_case_sensitive_regexp_count_scalar() {
    635
            let values = ["", "aabca", "abcabc", "abcAbcab", "abcabcabc"];
    636
            let regex = "abc";
    637
            let expected: Vec<i64> = vec![0, 1, 2, 1, 3];
    638
    639
            values.iter().enumerate().for_each(|(pos, &v)| {
    640
                // utf8
    641
                let v_sv = ScalarValue::Utf8(Some(v.to_string()));
    642
                let regex_sv = ScalarValue::Utf8(Some(regex.to_string()));
    643
                let expected = expected.get(pos).cloned();
    644
                let re = regexp_count_with_scalar_values(&[v_sv, regex_sv]);
    645
                match re {
    646
                    Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
    647
                        assert_eq!(v, expected, "regexp_count scalar test failed");
    648
                    }
    649
                    _ => panic!("Unexpected result"),
    650
                }
    651
    652
                // largeutf8
    653
                let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
    654
                let regex_sv = ScalarValue::LargeUtf8(Some(regex.to_string()));
    655
                let re = regexp_count_with_scalar_values(&[v_sv, regex_sv]);
    656
                match re {
    657
                    Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
    658
                        assert_eq!(v, expected, "regexp_count scalar test failed");
    659
                    }
    660
                    _ => panic!("Unexpected result"),
    661
                }
    662
    663
                // utf8view
    664
                let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
    665
                let regex_sv = ScalarValue::Utf8View(Some(regex.to_string()));
    666
                let re = regexp_count_with_scalar_values(&[v_sv, regex_sv]);
    667
                match re {
    668
                    Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
    669
                        assert_eq!(v, expected, "regexp_count scalar test failed");
    670
                    }
    671
                    _ => panic!("Unexpected result"),
    672
                }
    673
            });
    674
        }
    675
    676
        fn test_case_sensitive_regexp_count_scalar_start() {
    677
            let values = ["", "aabca", "abcabc", "abcAbcab", "abcabcabc"];
    678
            let regex = "abc";
    679
            let start = 2;
    680
            let expected: Vec<i64> = vec![0, 1, 1, 0, 2];
    681
    682
            values.iter().enumerate().for_each(|(pos, &v)| {
    683
                // utf8
    684
                let v_sv = ScalarValue::Utf8(Some(v.to_string()));
    685
                let regex_sv = ScalarValue::Utf8(Some(regex.to_string()));
    686
                let start_sv = ScalarValue::Int64(Some(start));
    687
                let expected = expected.get(pos).cloned();
    688
                let re = regexp_count_with_scalar_values(&[v_sv, regex_sv, start_sv.clone()]);
    689
                match re {
    690
                    Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
    691
                        assert_eq!(v, expected, "regexp_count scalar test failed");
    692
                    }
    693
                    _ => panic!("Unexpected result"),
    694
                }
    695
    696
                // largeutf8
    697
                let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
    698
                let regex_sv = ScalarValue::LargeUtf8(Some(regex.to_string()));
    699
                let re = regexp_count_with_scalar_values(&[v_sv, regex_sv, start_sv.clone()]);
    700
                match re {
    701
                    Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
    702
                        assert_eq!(v, expected, "regexp_count scalar test failed");
    703
                    }
    704
                    _ => panic!("Unexpected result"),
    705
                }
    706
    707
                // utf8view
    708
                let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
    709
                let regex_sv = ScalarValue::Utf8View(Some(regex.to_string()));
    710
                let re = regexp_count_with_scalar_values(&[v_sv, regex_sv, start_sv.clone()]);
    711
                match re {
    712
                    Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
    713
                        assert_eq!(v, expected, "regexp_count scalar test failed");
    714
                    }
    715
                    _ => panic!("Unexpected result"),
    716
                }
    717
            });
    718
        }
    719
    720
        fn test_case_insensitive_regexp_count_scalar_flags() {
    721
            let values = ["", "aabca", "abcabc", "abcAbcab", "abcabcabc"];
    722
            let regex = "abc";
    723
            let start = 1;
    724
            let flags = "i";
    725
            let expected: Vec<i64> = vec![0, 1, 2, 2, 3];
    726
    727
            values.iter().enumerate().for_each(|(pos, &v)| {
    728
                // utf8
    729
                let v_sv = ScalarValue::Utf8(Some(v.to_string()));
    730
                let regex_sv = ScalarValue::Utf8(Some(regex.to_string()));
    731
                let start_sv = ScalarValue::Int64(Some(start));
    732
                let flags_sv = ScalarValue::Utf8(Some(flags.to_string()));
    733
                let expected = expected.get(pos).cloned();
    734
    735
                let re = regexp_count_with_scalar_values(&[
    736
                    v_sv,
    737
                    regex_sv,
    738
                    start_sv.clone(),
    739
                    flags_sv.clone(),
    740
                ]);
    741
                match re {
    742
                    Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
    743
                        assert_eq!(v, expected, "regexp_count scalar test failed");
    744
                    }
    745
                    _ => panic!("Unexpected result"),
    746
                }
    747
    748
                // largeutf8
    749
                let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
    750
                let regex_sv = ScalarValue::LargeUtf8(Some(regex.to_string()));
    751
                let flags_sv = ScalarValue::LargeUtf8(Some(flags.to_string()));
    752
    753
                let re = regexp_count_with_scalar_values(&[
    754
                    v_sv,
    755
                    regex_sv,
    756
                    start_sv.clone(),
    757
                    flags_sv.clone(),
    758
                ]);
    759
                match re {
    760
                    Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
    761
                        assert_eq!(v, expected, "regexp_count scalar test failed");
    762
                    }
    763
                    _ => panic!("Unexpected result"),
    764
                }
    765
    766
                // utf8view
    767
                let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
    768
                let regex_sv = ScalarValue::Utf8View(Some(regex.to_string()));
    769
                let flags_sv = ScalarValue::Utf8View(Some(flags.to_string()));
    770
    771
                let re = regexp_count_with_scalar_values(&[
    772
                    v_sv,
    773
                    regex_sv,
    774
                    start_sv.clone(),
    775
                    flags_sv.clone(),
    776
                ]);
    777
                match re {
    778
                    Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
    779
                        assert_eq!(v, expected, "regexp_count scalar test failed");
    780
                    }
    781
                    _ => panic!("Unexpected result"),
    782
                }
    783
            });
    784
        }
    785
    786
        fn test_case_sensitive_regexp_count_array<A>()
    787
        where
    788
            A: From<Vec<&'static str>> + Array + 'static,
    789
        {
    790
            let values = A::from(vec!["", "aabca", "abcabc", "abcAbcab", "abcabcAbc"]);
    791
            let regex = A::from(vec!["", "abc", "a", "bc", "ab"]);
    792
    793
            let expected = Int64Array::from(vec![0, 1, 2, 2, 2]);
    794
    795
            let re = regexp_count_func(&[Arc::new(values), Arc::new(regex)]).unwrap();
    796
            assert_eq!(re.as_ref(), &expected);
    797
        }
    798
    799
        fn test_case_sensitive_regexp_count_array_start<A>()
    800
        where
    801
            A: From<Vec<&'static str>> + Array + 'static,
    802
        {
    803
            let values = A::from(vec!["", "aAbca", "abcabc", "abcAbcab", "abcabcAbc"]);
    804
            let regex = A::from(vec!["", "abc", "a", "bc", "ab"]);
    805
            let start = Int64Array::from(vec![1, 2, 3, 4, 5]);
    806
    807
            let expected = Int64Array::from(vec![0, 0, 1, 1, 0]);
    808
    809
            let re = regexp_count_func(&[Arc::new(values), Arc::new(regex), Arc::new(start)])
    810
                .unwrap();
    811
            assert_eq!(re.as_ref(), &expected);
    812
        }
    813
    814
        fn test_case_insensitive_regexp_count_array_flags<A>()
    815
        where
    816
            A: From<Vec<&'static str>> + Array + 'static,
    817
        {
    818
            let values = A::from(vec!["", "aAbca", "abcabc", "abcAbcab", "abcabcAbc"]);
    819
            let regex = A::from(vec!["", "abc", "a", "bc", "ab"]);
    820
            let start = Int64Array::from(vec![1]);
    821
            let flags = A::from(vec!["", "i", "", "", "i"]);
    822
    823
            let expected = Int64Array::from(vec![0, 1, 2, 2, 3]);
    824
    825
            let re = regexp_count_func(&[
    826
                Arc::new(values),
    827
                Arc::new(regex),
    828
                Arc::new(start),
    829
                Arc::new(flags),
    830
            ])
    831
            .unwrap();
    832
            assert_eq!(re.as_ref(), &expected);
    833
        }
    834
    835
        fn test_case_sensitive_regexp_count_start_scalar_complex() {
    836
            let values = ["", "aabca", "abcabc", "abcAbcab", "abcabcabc"];
    837
            let regex = ["", "abc", "a", "bc", "ab"];
    838
            let start = 5;
    839
            let flags = ["", "i", "", "", "i"];
    840
            let expected: Vec<i64> = vec![0, 0, 0, 1, 1];
    841
    842
            values.iter().enumerate().for_each(|(pos, &v)| {
    843
                // utf8
    844
                let v_sv = ScalarValue::Utf8(Some(v.to_string()));
    845
                let regex_sv = ScalarValue::Utf8(regex.get(pos).map(|s| (*s).to_string()));
    846
                let start_sv = ScalarValue::Int64(Some(start));
    847
                let flags_sv = ScalarValue::Utf8(flags.get(pos).map(|f| (*f).to_string()));
    848
                let expected = expected.get(pos).cloned();
    849
                let re = regexp_count_with_scalar_values(&[
    850
                    v_sv,
    851
                    regex_sv,
    852
                    start_sv.clone(),
    853
                    flags_sv.clone(),
    854
                ]);
    855
                match re {
    856
                    Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
    857
                        assert_eq!(v, expected, "regexp_count scalar test failed");
    858
                    }
    859
                    _ => panic!("Unexpected result"),
    860
                }
    861
    862
                // largeutf8
    863
                let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
    864
                let regex_sv =
    865
                    ScalarValue::LargeUtf8(regex.get(pos).map(|s| (*s).to_string()));
    866
                let flags_sv =
    867
                    ScalarValue::LargeUtf8(flags.get(pos).map(|f| (*f).to_string()));
    868
                let re = regexp_count_with_scalar_values(&[
    869
                    v_sv,
    870
                    regex_sv,
    871
                    start_sv.clone(),
    872
                    flags_sv.clone(),
    873
                ]);
    874
                match re {
    875
                    Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
    876
                        assert_eq!(v, expected, "regexp_count scalar test failed");
    877
                    }
    878
                    _ => panic!("Unexpected result"),
    879
                }
    880
    881
                // utf8view
    882
                let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
    883
                let regex_sv =
    884
                    ScalarValue::Utf8View(regex.get(pos).map(|s| (*s).to_string()));
    885
                let flags_sv =
    886
                    ScalarValue::Utf8View(flags.get(pos).map(|f| (*f).to_string()));
    887
                let re = regexp_count_with_scalar_values(&[
    888
                    v_sv,
    889
                    regex_sv,
    890
                    start_sv.clone(),
    891
                    flags_sv.clone(),
    892
                ]);
    893
                match re {
    894
                    Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
    895
                        assert_eq!(v, expected, "regexp_count scalar test failed");
    896
                    }
    897
                    _ => panic!("Unexpected result"),
    898
                }
    899
            });
    900
        }
    901
    902
        fn test_case_sensitive_regexp_count_array_complex<A>()
    903
        where
    904
            A: From<Vec<&'static str>> + Array + 'static,
    905
        {
    906
            let values = A::from(vec!["", "aAbca", "abcabc", "abcAbcab", "abcabcAbc"]);
    907
            let regex = A::from(vec!["", "abc", "a", "bc", "ab"]);
    908
            let start = Int64Array::from(vec![1, 2, 3, 4, 5]);
    909
            let flags = A::from(vec!["", "i", "", "", "i"]);
    910
    911
            let expected = Int64Array::from(vec![0, 1, 1, 1, 1]);
    912
    913
            let re = regexp_count_func(&[
    914
                Arc::new(values),
    915
                Arc::new(regex),
    916
                Arc::new(start),
    917
                Arc::new(flags),
    918
            ])
    919
            .unwrap();
    920
            assert_eq!(re.as_ref(), &expected);
    921
        }
    922
    923
        fn test_case_regexp_count_cache_check<A>()
    924
        where
    925
            A: From<Vec<&'static str>> + Array + 'static,
    926
        {
    927
            let values = A::from(vec!["aaa", "Aaa", "aaa"]);
    928
            let regex = A::from(vec!["aaa", "aaa", "aaa"]);
    929
            let start = Int64Array::from(vec![1, 1, 1]);
    930
            let flags = A::from(vec!["", "i", ""]);
    931
    932
            let expected = Int64Array::from(vec![1, 1, 1]);
    933
    934
            let re = regexp_count_func(&[
    935
                Arc::new(values),
    936
                Arc::new(regex),
    937
                Arc::new(start),
    938
                Arc::new(flags),
    939
            ])
    940
            .unwrap();
    941
            assert_eq!(re.as_ref(), &expected);
    942
        }
    943
    }