datafusion_functions/regex/
regexpcount.rs

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
18use crate::regex::{compile_and_cache_regex, compile_regex};
19use arrow::array::{Array, ArrayRef, AsArray, Datum, Int64Array, StringArrayType};
20use arrow::datatypes::{DataType, Int64Type};
21use arrow::datatypes::{
22    DataType::Int64, DataType::LargeUtf8, DataType::Utf8, DataType::Utf8View,
23};
24use arrow::error::ArrowError;
25use datafusion_common::{exec_err, internal_err, Result, ScalarValue};
26use datafusion_expr::{
27    ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignature::Exact,
28    TypeSignature::Uniform, Volatility,
29};
30use datafusion_macros::user_doc;
31use itertools::izip;
32use regex::Regex;
33use std::collections::HashMap;
34use 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)]
65pub struct RegexpCountFunc {
66    signature: Signature,
67}
68
69impl Default for RegexpCountFunc {
70    fn default() -> Self {
71        Self::new()
72    }
73}
74
75impl RegexpCountFunc {
76    pub fn new() -> Self {
77        Self {
78            signature: Signature::one_of(
79                vec![
80                    Uniform(2, vec![Utf8View, LargeUtf8, Utf8]),
81                    Exact(vec![Utf8View, Utf8View, Int64]),
82                    Exact(vec![LargeUtf8, LargeUtf8, Int64]),
83                    Exact(vec![Utf8, Utf8, Int64]),
84                    Exact(vec![Utf8View, Utf8View, Int64, Utf8View]),
85                    Exact(vec![LargeUtf8, LargeUtf8, Int64, LargeUtf8]),
86                    Exact(vec![Utf8, Utf8, Int64, Utf8]),
87                ],
88                Volatility::Immutable,
89            ),
90        }
91    }
92}
93
94impl ScalarUDFImpl for RegexpCountFunc {
95    fn as_any(&self) -> &dyn std::any::Any {
96        self
97    }
98
99    fn name(&self) -> &str {
100        "regexp_count"
101    }
102
103    fn signature(&self) -> &Signature {
104        &self.signature
105    }
106
107    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
108        Ok(Int64)
109    }
110
111    fn invoke_with_args(
112        &self,
113        args: datafusion_expr::ScalarFunctionArgs,
114    ) -> Result<ColumnarValue> {
115        let args = &args.args;
116
117        let len = args
118            .iter()
119            .fold(Option::<usize>::None, |acc, arg| match arg {
120                ColumnarValue::Scalar(_) => acc,
121                ColumnarValue::Array(a) => Some(a.len()),
122            });
123
124        let is_scalar = len.is_none();
125        let inferred_length = len.unwrap_or(1);
126        let args = args
127            .iter()
128            .map(|arg| arg.to_array(inferred_length))
129            .collect::<Result<Vec<_>>>()?;
130
131        let result = regexp_count_func(&args);
132        if is_scalar {
133            // If all inputs are scalar, keeps output as scalar
134            let result = result.and_then(|arr| ScalarValue::try_from_array(&arr, 0));
135            result.map(ColumnarValue::Scalar)
136        } else {
137            result.map(ColumnarValue::Array)
138        }
139    }
140
141    fn documentation(&self) -> Option<&Documentation> {
142        self.doc()
143    }
144}
145
146pub fn regexp_count_func(args: &[ArrayRef]) -> Result<ArrayRef> {
147    let args_len = args.len();
148    if !(2..=4).contains(&args_len) {
149        return exec_err!("regexp_count was called with {args_len} arguments. It requires at least 2 and at most 4.");
150    }
151
152    let values = &args[0];
153    match values.data_type() {
154        Utf8 | LargeUtf8 | Utf8View => (),
155        other => {
156            return internal_err!(
157                "Unsupported data type {other:?} for function regexp_count"
158            );
159        }
160    }
161
162    regexp_count(
163        values,
164        &args[1],
165        if args_len > 2 { Some(&args[2]) } else { None },
166        if args_len > 3 { Some(&args[3]) } else { None },
167    )
168    .map_err(|e| e.into())
169}
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.
186pub fn regexp_count(
187    values: &dyn Array,
188    regex_array: &dyn Datum,
189    start_array: Option<&dyn Datum>,
190    flags_array: Option<&dyn Datum>,
191) -> Result<ArrayRef, ArrowError> {
192    let (regex_array, is_regex_scalar) = regex_array.get();
193    let (start_array, is_start_scalar) = start_array.map_or((None, true), |start| {
194        let (start, is_start_scalar) = start.get();
195        (Some(start), is_start_scalar)
196    });
197    let (flags_array, is_flags_scalar) = flags_array.map_or((None, true), |flags| {
198        let (flags, is_flags_scalar) = flags.get();
199        (Some(flags), is_flags_scalar)
200    });
201
202    match (values.data_type(), regex_array.data_type(), flags_array) {
203        (Utf8, Utf8, None) => regexp_count_inner(
204            values.as_string::<i32>(),
205            regex_array.as_string::<i32>(),
206            is_regex_scalar,
207            start_array.map(|start| start.as_primitive::<Int64Type>()),
208            is_start_scalar,
209            None,
210            is_flags_scalar,
211        ),
212        (Utf8, Utf8, Some(flags_array)) if *flags_array.data_type() == Utf8 => regexp_count_inner(
213            values.as_string::<i32>(),
214            regex_array.as_string::<i32>(),
215            is_regex_scalar,
216            start_array.map(|start| start.as_primitive::<Int64Type>()),
217            is_start_scalar,
218            Some(flags_array.as_string::<i32>()),
219            is_flags_scalar,
220        ),
221        (LargeUtf8, LargeUtf8, None) => regexp_count_inner(
222            values.as_string::<i64>(),
223            regex_array.as_string::<i64>(),
224            is_regex_scalar,
225            start_array.map(|start| start.as_primitive::<Int64Type>()),
226            is_start_scalar,
227            None,
228            is_flags_scalar,
229        ),
230        (LargeUtf8, LargeUtf8, Some(flags_array)) if *flags_array.data_type() == LargeUtf8 => regexp_count_inner(
231            values.as_string::<i64>(),
232            regex_array.as_string::<i64>(),
233            is_regex_scalar,
234            start_array.map(|start| start.as_primitive::<Int64Type>()),
235            is_start_scalar,
236            Some(flags_array.as_string::<i64>()),
237            is_flags_scalar,
238        ),
239        (Utf8View, Utf8View, None) => regexp_count_inner(
240            values.as_string_view(),
241            regex_array.as_string_view(),
242            is_regex_scalar,
243            start_array.map(|start| start.as_primitive::<Int64Type>()),
244            is_start_scalar,
245            None,
246            is_flags_scalar,
247        ),
248        (Utf8View, Utf8View, Some(flags_array)) if *flags_array.data_type() == Utf8View => regexp_count_inner(
249            values.as_string_view(),
250            regex_array.as_string_view(),
251            is_regex_scalar,
252            start_array.map(|start| start.as_primitive::<Int64Type>()),
253            is_start_scalar,
254            Some(flags_array.as_string_view()),
255            is_flags_scalar,
256        ),
257        _ => Err(ArrowError::ComputeError(
258            "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        )),
260    }
261}
262
263pub fn regexp_count_inner<'a, S>(
264    values: S,
265    regex_array: S,
266    is_regex_scalar: bool,
267    start_array: Option<&Int64Array>,
268    is_start_scalar: bool,
269    flags_array: Option<S>,
270    is_flags_scalar: bool,
271) -> Result<ArrayRef, ArrowError>
272where
273    S: StringArrayType<'a>,
274{
275    let (regex_scalar, is_regex_scalar) = if is_regex_scalar || regex_array.len() == 1 {
276        (Some(regex_array.value(0)), true)
277    } else {
278        (None, false)
279    };
280
281    let (start_array, start_scalar, is_start_scalar) =
282        if let Some(start_array) = start_array {
283            if is_start_scalar || start_array.len() == 1 {
284                (None, Some(start_array.value(0)), true)
285            } else {
286                (Some(start_array), None, false)
287            }
288        } else {
289            (None, Some(1), true)
290        };
291
292    let (flags_array, flags_scalar, is_flags_scalar) =
293        if let Some(flags_array) = flags_array {
294            if is_flags_scalar || flags_array.len() == 1 {
295                (None, Some(flags_array.value(0)), true)
296            } else {
297                (Some(flags_array), None, false)
298            }
299        } else {
300            (None, None, true)
301        };
302
303    let mut regex_cache = HashMap::new();
304
305    match (is_regex_scalar, is_start_scalar, is_flags_scalar) {
306        (true, true, true) => {
307            let regex = match regex_scalar {
308                None | Some("") => {
309                    return Ok(Arc::new(Int64Array::from(vec![0; values.len()])))
310                }
311                Some(regex) => regex,
312            };
313
314            let pattern = compile_regex(regex, flags_scalar)?;
315
316            Ok(Arc::new(
317                values
318                    .iter()
319                    .map(|value| count_matches(value, &pattern, start_scalar))
320                    .collect::<Result<Int64Array, ArrowError>>()?,
321            ))
322        }
323        (true, true, false) => {
324            let regex = match regex_scalar {
325                None | Some("") => {
326                    return Ok(Arc::new(Int64Array::from(vec![0; values.len()])))
327                }
328                Some(regex) => regex,
329            };
330
331            let flags_array = flags_array.unwrap();
332            if values.len() != flags_array.len() {
333                return Err(ArrowError::ComputeError(format!(
334                    "flags_array must be the same length as values array; got {} and {}",
335                    flags_array.len(),
336                    values.len(),
337                )));
338            }
339
340            Ok(Arc::new(
341                values
342                    .iter()
343                    .zip(flags_array.iter())
344                    .map(|(value, flags)| {
345                        let pattern =
346                            compile_and_cache_regex(regex, flags, &mut regex_cache)?;
347                        count_matches(value, pattern, start_scalar)
348                    })
349                    .collect::<Result<Int64Array, ArrowError>>()?,
350            ))
351        }
352        (true, false, true) => {
353            let regex = match regex_scalar {
354                None | Some("") => {
355                    return Ok(Arc::new(Int64Array::from(vec![0; values.len()])))
356                }
357                Some(regex) => regex,
358            };
359
360            let pattern = compile_regex(regex, flags_scalar)?;
361
362            let start_array = start_array.unwrap();
363
364            Ok(Arc::new(
365                values
366                    .iter()
367                    .zip(start_array.iter())
368                    .map(|(value, start)| count_matches(value, &pattern, start))
369                    .collect::<Result<Int64Array, ArrowError>>()?,
370            ))
371        }
372        (true, false, false) => {
373            let regex = match regex_scalar {
374                None | Some("") => {
375                    return Ok(Arc::new(Int64Array::from(vec![0; values.len()])))
376                }
377                Some(regex) => regex,
378            };
379
380            let flags_array = flags_array.unwrap();
381            if values.len() != flags_array.len() {
382                return Err(ArrowError::ComputeError(format!(
383                    "flags_array must be the same length as values array; got {} and {}",
384                    flags_array.len(),
385                    values.len(),
386                )));
387            }
388
389            Ok(Arc::new(
390                izip!(
391                    values.iter(),
392                    start_array.unwrap().iter(),
393                    flags_array.iter()
394                )
395                .map(|(value, start, flags)| {
396                    let pattern =
397                        compile_and_cache_regex(regex, flags, &mut regex_cache)?;
398
399                    count_matches(value, pattern, start)
400                })
401                .collect::<Result<Int64Array, ArrowError>>()?,
402            ))
403        }
404        (false, true, true) => {
405            if values.len() != regex_array.len() {
406                return Err(ArrowError::ComputeError(format!(
407                    "regex_array must be the same length as values array; got {} and {}",
408                    regex_array.len(),
409                    values.len(),
410                )));
411            }
412
413            Ok(Arc::new(
414                values
415                    .iter()
416                    .zip(regex_array.iter())
417                    .map(|(value, regex)| {
418                        let regex = match regex {
419                            None | Some("") => return Ok(0),
420                            Some(regex) => regex,
421                        };
422
423                        let pattern = compile_and_cache_regex(
424                            regex,
425                            flags_scalar,
426                            &mut regex_cache,
427                        )?;
428                        count_matches(value, pattern, start_scalar)
429                    })
430                    .collect::<Result<Int64Array, ArrowError>>()?,
431            ))
432        }
433        (false, true, false) => {
434            if values.len() != regex_array.len() {
435                return Err(ArrowError::ComputeError(format!(
436                    "regex_array must be the same length as values array; got {} and {}",
437                    regex_array.len(),
438                    values.len(),
439                )));
440            }
441
442            let flags_array = flags_array.unwrap();
443            if values.len() != flags_array.len() {
444                return Err(ArrowError::ComputeError(format!(
445                    "flags_array must be the same length as values array; got {} and {}",
446                    flags_array.len(),
447                    values.len(),
448                )));
449            }
450
451            Ok(Arc::new(
452                izip!(values.iter(), regex_array.iter(), flags_array.iter())
453                    .map(|(value, regex, flags)| {
454                        let regex = match regex {
455                            None | Some("") => return Ok(0),
456                            Some(regex) => regex,
457                        };
458
459                        let pattern =
460                            compile_and_cache_regex(regex, flags, &mut regex_cache)?;
461
462                        count_matches(value, pattern, start_scalar)
463                    })
464                    .collect::<Result<Int64Array, ArrowError>>()?,
465            ))
466        }
467        (false, false, true) => {
468            if values.len() != regex_array.len() {
469                return Err(ArrowError::ComputeError(format!(
470                    "regex_array must be the same length as values array; got {} and {}",
471                    regex_array.len(),
472                    values.len(),
473                )));
474            }
475
476            let start_array = start_array.unwrap();
477            if values.len() != start_array.len() {
478                return Err(ArrowError::ComputeError(format!(
479                    "start_array must be the same length as values array; got {} and {}",
480                    start_array.len(),
481                    values.len(),
482                )));
483            }
484
485            Ok(Arc::new(
486                izip!(values.iter(), regex_array.iter(), start_array.iter())
487                    .map(|(value, regex, start)| {
488                        let regex = match regex {
489                            None | Some("") => return Ok(0),
490                            Some(regex) => regex,
491                        };
492
493                        let pattern = compile_and_cache_regex(
494                            regex,
495                            flags_scalar,
496                            &mut regex_cache,
497                        )?;
498                        count_matches(value, pattern, start)
499                    })
500                    .collect::<Result<Int64Array, ArrowError>>()?,
501            ))
502        }
503        (false, false, false) => {
504            if values.len() != regex_array.len() {
505                return Err(ArrowError::ComputeError(format!(
506                    "regex_array must be the same length as values array; got {} and {}",
507                    regex_array.len(),
508                    values.len(),
509                )));
510            }
511
512            let start_array = start_array.unwrap();
513            if values.len() != start_array.len() {
514                return Err(ArrowError::ComputeError(format!(
515                    "start_array must be the same length as values array; got {} and {}",
516                    start_array.len(),
517                    values.len(),
518                )));
519            }
520
521            let flags_array = flags_array.unwrap();
522            if values.len() != flags_array.len() {
523                return Err(ArrowError::ComputeError(format!(
524                    "flags_array must be the same length as values array; got {} and {}",
525                    flags_array.len(),
526                    values.len(),
527                )));
528            }
529
530            Ok(Arc::new(
531                izip!(
532                    values.iter(),
533                    regex_array.iter(),
534                    start_array.iter(),
535                    flags_array.iter()
536                )
537                .map(|(value, regex, start, flags)| {
538                    let regex = match regex {
539                        None | Some("") => return Ok(0),
540                        Some(regex) => regex,
541                    };
542
543                    let pattern =
544                        compile_and_cache_regex(regex, flags, &mut regex_cache)?;
545                    count_matches(value, pattern, start)
546                })
547                .collect::<Result<Int64Array, ArrowError>>()?,
548            ))
549        }
550    }
551}
552
553fn count_matches(
554    value: Option<&str>,
555    pattern: &Regex,
556    start: Option<i64>,
557) -> Result<i64, ArrowError> {
558    let value = match value {
559        None | Some("") => return Ok(0),
560        Some(value) => value,
561    };
562
563    if let Some(start) = start {
564        if start < 1 {
565            return Err(ArrowError::ComputeError(
566                "regexp_count() requires start to be 1 based".to_string(),
567            ));
568        }
569
570        let find_slice = value.chars().skip(start as usize - 1).collect::<String>();
571        let count = pattern.find_iter(find_slice.as_str()).count();
572        Ok(count as i64)
573    } else {
574        let count = pattern.find_iter(value).count();
575        Ok(count as i64)
576    }
577}
578
579#[cfg(test)]
580mod 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 = ScalarValue::LargeUtf8(regex.get(pos).map(|s| s.to_string()));
865            let flags_sv = ScalarValue::LargeUtf8(flags.get(pos).map(|f| f.to_string()));
866            let re = regexp_count_with_scalar_values(&[
867                v_sv,
868                regex_sv,
869                start_sv.clone(),
870                flags_sv.clone(),
871            ]);
872            match re {
873                Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
874                    assert_eq!(v, expected, "regexp_count scalar test failed");
875                }
876                _ => panic!("Unexpected result"),
877            }
878
879            // utf8view
880            let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
881            let regex_sv = ScalarValue::Utf8View(regex.get(pos).map(|s| s.to_string()));
882            let flags_sv = ScalarValue::Utf8View(flags.get(pos).map(|f| f.to_string()));
883            let re = regexp_count_with_scalar_values(&[
884                v_sv,
885                regex_sv,
886                start_sv.clone(),
887                flags_sv.clone(),
888            ]);
889            match re {
890                Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
891                    assert_eq!(v, expected, "regexp_count scalar test failed");
892                }
893                _ => panic!("Unexpected result"),
894            }
895        });
896    }
897
898    fn test_case_sensitive_regexp_count_array_complex<A>()
899    where
900        A: From<Vec<&'static str>> + Array + 'static,
901    {
902        let values = A::from(vec!["", "aAbca", "abcabc", "abcAbcab", "abcabcAbc"]);
903        let regex = A::from(vec!["", "abc", "a", "bc", "ab"]);
904        let start = Int64Array::from(vec![1, 2, 3, 4, 5]);
905        let flags = A::from(vec!["", "i", "", "", "i"]);
906
907        let expected = Int64Array::from(vec![0, 1, 1, 1, 1]);
908
909        let re = regexp_count_func(&[
910            Arc::new(values),
911            Arc::new(regex),
912            Arc::new(start),
913            Arc::new(flags),
914        ])
915        .unwrap();
916        assert_eq!(re.as_ref(), &expected);
917    }
918
919    fn test_case_regexp_count_cache_check<A>()
920    where
921        A: From<Vec<&'static str>> + Array + 'static,
922    {
923        let values = A::from(vec!["aaa", "Aaa", "aaa"]);
924        let regex = A::from(vec!["aaa", "aaa", "aaa"]);
925        let start = Int64Array::from(vec![1, 1, 1]);
926        let flags = A::from(vec!["", "i", ""]);
927
928        let expected = Int64Array::from(vec![1, 1, 1]);
929
930        let re = regexp_count_func(&[
931            Arc::new(values),
932            Arc::new(regex),
933            Arc::new(start),
934            Arc::new(flags),
935        ])
936        .unwrap();
937        assert_eq!(re.as_ref(), &expected);
938    }
939}