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            lambdas: None,
632        })
633    }
634
635    fn test_case_sensitive_regexp_count_scalar() {
636        let values = ["", "aabca", "abcabc", "abcAbcab", "abcabcabc"];
637        let regex = "abc";
638        let expected: Vec<i64> = vec![0, 1, 2, 1, 3];
639
640        values.iter().enumerate().for_each(|(pos, &v)| {
641            // utf8
642            let v_sv = ScalarValue::Utf8(Some(v.to_string()));
643            let regex_sv = ScalarValue::Utf8(Some(regex.to_string()));
644            let expected = expected.get(pos).cloned();
645            let re = regexp_count_with_scalar_values(&[v_sv, regex_sv]);
646            match re {
647                Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
648                    assert_eq!(v, expected, "regexp_count scalar test failed");
649                }
650                _ => panic!("Unexpected result"),
651            }
652
653            // largeutf8
654            let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
655            let regex_sv = ScalarValue::LargeUtf8(Some(regex.to_string()));
656            let re = regexp_count_with_scalar_values(&[v_sv, regex_sv]);
657            match re {
658                Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
659                    assert_eq!(v, expected, "regexp_count scalar test failed");
660                }
661                _ => panic!("Unexpected result"),
662            }
663
664            // utf8view
665            let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
666            let regex_sv = ScalarValue::Utf8View(Some(regex.to_string()));
667            let re = regexp_count_with_scalar_values(&[v_sv, regex_sv]);
668            match re {
669                Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
670                    assert_eq!(v, expected, "regexp_count scalar test failed");
671                }
672                _ => panic!("Unexpected result"),
673            }
674        });
675    }
676
677    fn test_case_sensitive_regexp_count_scalar_start() {
678        let values = ["", "aabca", "abcabc", "abcAbcab", "abcabcabc"];
679        let regex = "abc";
680        let start = 2;
681        let expected: Vec<i64> = vec![0, 1, 1, 0, 2];
682
683        values.iter().enumerate().for_each(|(pos, &v)| {
684            // utf8
685            let v_sv = ScalarValue::Utf8(Some(v.to_string()));
686            let regex_sv = ScalarValue::Utf8(Some(regex.to_string()));
687            let start_sv = ScalarValue::Int64(Some(start));
688            let expected = expected.get(pos).cloned();
689            let re = regexp_count_with_scalar_values(&[v_sv, regex_sv, start_sv.clone()]);
690            match re {
691                Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
692                    assert_eq!(v, expected, "regexp_count scalar test failed");
693                }
694                _ => panic!("Unexpected result"),
695            }
696
697            // largeutf8
698            let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
699            let regex_sv = ScalarValue::LargeUtf8(Some(regex.to_string()));
700            let re = regexp_count_with_scalar_values(&[v_sv, regex_sv, start_sv.clone()]);
701            match re {
702                Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
703                    assert_eq!(v, expected, "regexp_count scalar test failed");
704                }
705                _ => panic!("Unexpected result"),
706            }
707
708            // utf8view
709            let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
710            let regex_sv = ScalarValue::Utf8View(Some(regex.to_string()));
711            let re = regexp_count_with_scalar_values(&[v_sv, regex_sv, start_sv.clone()]);
712            match re {
713                Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
714                    assert_eq!(v, expected, "regexp_count scalar test failed");
715                }
716                _ => panic!("Unexpected result"),
717            }
718        });
719    }
720
721    fn test_case_insensitive_regexp_count_scalar_flags() {
722        let values = ["", "aabca", "abcabc", "abcAbcab", "abcabcabc"];
723        let regex = "abc";
724        let start = 1;
725        let flags = "i";
726        let expected: Vec<i64> = vec![0, 1, 2, 2, 3];
727
728        values.iter().enumerate().for_each(|(pos, &v)| {
729            // utf8
730            let v_sv = ScalarValue::Utf8(Some(v.to_string()));
731            let regex_sv = ScalarValue::Utf8(Some(regex.to_string()));
732            let start_sv = ScalarValue::Int64(Some(start));
733            let flags_sv = ScalarValue::Utf8(Some(flags.to_string()));
734            let expected = expected.get(pos).cloned();
735
736            let re = regexp_count_with_scalar_values(&[
737                v_sv,
738                regex_sv,
739                start_sv.clone(),
740                flags_sv.clone(),
741            ]);
742            match re {
743                Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
744                    assert_eq!(v, expected, "regexp_count scalar test failed");
745                }
746                _ => panic!("Unexpected result"),
747            }
748
749            // largeutf8
750            let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
751            let regex_sv = ScalarValue::LargeUtf8(Some(regex.to_string()));
752            let flags_sv = ScalarValue::LargeUtf8(Some(flags.to_string()));
753
754            let re = regexp_count_with_scalar_values(&[
755                v_sv,
756                regex_sv,
757                start_sv.clone(),
758                flags_sv.clone(),
759            ]);
760            match re {
761                Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
762                    assert_eq!(v, expected, "regexp_count scalar test failed");
763                }
764                _ => panic!("Unexpected result"),
765            }
766
767            // utf8view
768            let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
769            let regex_sv = ScalarValue::Utf8View(Some(regex.to_string()));
770            let flags_sv = ScalarValue::Utf8View(Some(flags.to_string()));
771
772            let re = regexp_count_with_scalar_values(&[
773                v_sv,
774                regex_sv,
775                start_sv.clone(),
776                flags_sv.clone(),
777            ]);
778            match re {
779                Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
780                    assert_eq!(v, expected, "regexp_count scalar test failed");
781                }
782                _ => panic!("Unexpected result"),
783            }
784        });
785    }
786
787    fn test_case_sensitive_regexp_count_array<A>()
788    where
789        A: From<Vec<&'static str>> + Array + 'static,
790    {
791        let values = A::from(vec!["", "aabca", "abcabc", "abcAbcab", "abcabcAbc"]);
792        let regex = A::from(vec!["", "abc", "a", "bc", "ab"]);
793
794        let expected = Int64Array::from(vec![0, 1, 2, 2, 2]);
795
796        let re = regexp_count_func(&[Arc::new(values), Arc::new(regex)]).unwrap();
797        assert_eq!(re.as_ref(), &expected);
798    }
799
800    fn test_case_sensitive_regexp_count_array_start<A>()
801    where
802        A: From<Vec<&'static str>> + Array + 'static,
803    {
804        let values = A::from(vec!["", "aAbca", "abcabc", "abcAbcab", "abcabcAbc"]);
805        let regex = A::from(vec!["", "abc", "a", "bc", "ab"]);
806        let start = Int64Array::from(vec![1, 2, 3, 4, 5]);
807
808        let expected = Int64Array::from(vec![0, 0, 1, 1, 0]);
809
810        let re = regexp_count_func(&[Arc::new(values), Arc::new(regex), Arc::new(start)])
811            .unwrap();
812        assert_eq!(re.as_ref(), &expected);
813    }
814
815    fn test_case_insensitive_regexp_count_array_flags<A>()
816    where
817        A: From<Vec<&'static str>> + Array + 'static,
818    {
819        let values = A::from(vec!["", "aAbca", "abcabc", "abcAbcab", "abcabcAbc"]);
820        let regex = A::from(vec!["", "abc", "a", "bc", "ab"]);
821        let start = Int64Array::from(vec![1]);
822        let flags = A::from(vec!["", "i", "", "", "i"]);
823
824        let expected = Int64Array::from(vec![0, 1, 2, 2, 3]);
825
826        let re = regexp_count_func(&[
827            Arc::new(values),
828            Arc::new(regex),
829            Arc::new(start),
830            Arc::new(flags),
831        ])
832        .unwrap();
833        assert_eq!(re.as_ref(), &expected);
834    }
835
836    fn test_case_sensitive_regexp_count_start_scalar_complex() {
837        let values = ["", "aabca", "abcabc", "abcAbcab", "abcabcabc"];
838        let regex = ["", "abc", "a", "bc", "ab"];
839        let start = 5;
840        let flags = ["", "i", "", "", "i"];
841        let expected: Vec<i64> = vec![0, 0, 0, 1, 1];
842
843        values.iter().enumerate().for_each(|(pos, &v)| {
844            // utf8
845            let v_sv = ScalarValue::Utf8(Some(v.to_string()));
846            let regex_sv = ScalarValue::Utf8(regex.get(pos).map(|s| (*s).to_string()));
847            let start_sv = ScalarValue::Int64(Some(start));
848            let flags_sv = ScalarValue::Utf8(flags.get(pos).map(|f| (*f).to_string()));
849            let expected = expected.get(pos).cloned();
850            let re = regexp_count_with_scalar_values(&[
851                v_sv,
852                regex_sv,
853                start_sv.clone(),
854                flags_sv.clone(),
855            ]);
856            match re {
857                Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
858                    assert_eq!(v, expected, "regexp_count scalar test failed");
859                }
860                _ => panic!("Unexpected result"),
861            }
862
863            // largeutf8
864            let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
865            let regex_sv =
866                ScalarValue::LargeUtf8(regex.get(pos).map(|s| (*s).to_string()));
867            let flags_sv =
868                ScalarValue::LargeUtf8(flags.get(pos).map(|f| (*f).to_string()));
869            let re = regexp_count_with_scalar_values(&[
870                v_sv,
871                regex_sv,
872                start_sv.clone(),
873                flags_sv.clone(),
874            ]);
875            match re {
876                Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
877                    assert_eq!(v, expected, "regexp_count scalar test failed");
878                }
879                _ => panic!("Unexpected result"),
880            }
881
882            // utf8view
883            let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
884            let regex_sv =
885                ScalarValue::Utf8View(regex.get(pos).map(|s| (*s).to_string()));
886            let flags_sv =
887                ScalarValue::Utf8View(flags.get(pos).map(|f| (*f).to_string()));
888            let re = regexp_count_with_scalar_values(&[
889                v_sv,
890                regex_sv,
891                start_sv.clone(),
892                flags_sv.clone(),
893            ]);
894            match re {
895                Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
896                    assert_eq!(v, expected, "regexp_count scalar test failed");
897                }
898                _ => panic!("Unexpected result"),
899            }
900        });
901    }
902
903    fn test_case_sensitive_regexp_count_array_complex<A>()
904    where
905        A: From<Vec<&'static str>> + Array + 'static,
906    {
907        let values = A::from(vec!["", "aAbca", "abcabc", "abcAbcab", "abcabcAbc"]);
908        let regex = A::from(vec!["", "abc", "a", "bc", "ab"]);
909        let start = Int64Array::from(vec![1, 2, 3, 4, 5]);
910        let flags = A::from(vec!["", "i", "", "", "i"]);
911
912        let expected = Int64Array::from(vec![0, 1, 1, 1, 1]);
913
914        let re = regexp_count_func(&[
915            Arc::new(values),
916            Arc::new(regex),
917            Arc::new(start),
918            Arc::new(flags),
919        ])
920        .unwrap();
921        assert_eq!(re.as_ref(), &expected);
922    }
923
924    fn test_case_regexp_count_cache_check<A>()
925    where
926        A: From<Vec<&'static str>> + Array + 'static,
927    {
928        let values = A::from(vec!["aaa", "Aaa", "aaa"]);
929        let regex = A::from(vec!["aaa", "aaa", "aaa"]);
930        let start = Int64Array::from(vec![1, 1, 1]);
931        let flags = A::from(vec!["", "i", ""]);
932
933        let expected = Int64Array::from(vec![1, 1, 1]);
934
935        let re = regexp_count_func(&[
936            Arc::new(values),
937            Arc::new(regex),
938            Arc::new(start),
939            Arc::new(flags),
940        ])
941        .unwrap();
942        assert_eq!(re.as_ref(), &expected);
943    }
944}