1use arrow::array::{Array, ArrayRef, AsArray, Datum, Int64Array, StringArrayType};
19use arrow::datatypes::{DataType, Int64Type};
20use arrow::datatypes::{
21 DataType::Int64, DataType::LargeUtf8, DataType::Utf8, DataType::Utf8View,
22};
23use arrow::error::ArrowError;
24use datafusion_common::{exec_err, internal_err, Result, ScalarValue};
25use datafusion_expr::{
26 ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignature::Exact,
27 TypeSignature::Uniform, Volatility,
28};
29use datafusion_macros::user_doc;
30use itertools::izip;
31use regex::Regex;
32use std::collections::hash_map::Entry;
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)]
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 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
171pub 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 compile_and_cache_regex<'strings, 'cache>(
554 regex: &'strings str,
555 flags: Option<&'strings str>,
556 regex_cache: &'cache mut HashMap<(&'strings str, Option<&'strings str>), Regex>,
557) -> Result<&'cache Regex, ArrowError>
558where
559 'strings: 'cache,
560{
561 let result = match regex_cache.entry((regex, flags)) {
562 Entry::Occupied(occupied_entry) => occupied_entry.into_mut(),
563 Entry::Vacant(vacant_entry) => {
564 let compiled = compile_regex(regex, flags)?;
565 vacant_entry.insert(compiled)
566 }
567 };
568 Ok(result)
569}
570
571fn compile_regex(regex: &str, flags: Option<&str>) -> Result<Regex, ArrowError> {
572 let pattern = match flags {
573 None | Some("") => regex.to_string(),
574 Some(flags) => {
575 if flags.contains("g") {
576 return Err(ArrowError::ComputeError(
577 "regexp_count() does not support global flag".to_string(),
578 ));
579 }
580 format!("(?{flags}){regex}")
581 }
582 };
583
584 Regex::new(&pattern).map_err(|_| {
585 ArrowError::ComputeError(format!("Regular expression did not compile: {pattern}"))
586 })
587}
588
589fn count_matches(
590 value: Option<&str>,
591 pattern: &Regex,
592 start: Option<i64>,
593) -> Result<i64, ArrowError> {
594 let value = match value {
595 None | Some("") => return Ok(0),
596 Some(value) => value,
597 };
598
599 if let Some(start) = start {
600 if start < 1 {
601 return Err(ArrowError::ComputeError(
602 "regexp_count() requires start to be 1 based".to_string(),
603 ));
604 }
605
606 let find_slice = value.chars().skip(start as usize - 1).collect::<String>();
607 let count = pattern.find_iter(find_slice.as_str()).count();
608 Ok(count as i64)
609 } else {
610 let count = pattern.find_iter(value).count();
611 Ok(count as i64)
612 }
613}
614
615#[cfg(test)]
616mod tests {
617 use super::*;
618 use arrow::array::{GenericStringArray, StringViewArray};
619 use arrow::datatypes::Field;
620 use datafusion_expr::ScalarFunctionArgs;
621
622 #[test]
623 fn test_regexp_count() {
624 test_case_sensitive_regexp_count_scalar();
625 test_case_sensitive_regexp_count_scalar_start();
626 test_case_insensitive_regexp_count_scalar_flags();
627 test_case_sensitive_regexp_count_start_scalar_complex();
628
629 test_case_sensitive_regexp_count_array::<GenericStringArray<i32>>();
630 test_case_sensitive_regexp_count_array::<GenericStringArray<i64>>();
631 test_case_sensitive_regexp_count_array::<StringViewArray>();
632
633 test_case_sensitive_regexp_count_array_start::<GenericStringArray<i32>>();
634 test_case_sensitive_regexp_count_array_start::<GenericStringArray<i64>>();
635 test_case_sensitive_regexp_count_array_start::<StringViewArray>();
636
637 test_case_insensitive_regexp_count_array_flags::<GenericStringArray<i32>>();
638 test_case_insensitive_regexp_count_array_flags::<GenericStringArray<i64>>();
639 test_case_insensitive_regexp_count_array_flags::<StringViewArray>();
640
641 test_case_sensitive_regexp_count_array_complex::<GenericStringArray<i32>>();
642 test_case_sensitive_regexp_count_array_complex::<GenericStringArray<i64>>();
643 test_case_sensitive_regexp_count_array_complex::<StringViewArray>();
644
645 test_case_regexp_count_cache_check::<GenericStringArray<i32>>();
646 }
647
648 fn regexp_count_with_scalar_values(args: &[ScalarValue]) -> Result<ColumnarValue> {
649 let args_values = args
650 .iter()
651 .map(|sv| ColumnarValue::Scalar(sv.clone()))
652 .collect();
653
654 let arg_fields = args
655 .iter()
656 .enumerate()
657 .map(|(idx, a)| Field::new(format!("arg_{idx}"), a.data_type(), true).into())
658 .collect::<Vec<_>>();
659
660 RegexpCountFunc::new().invoke_with_args(ScalarFunctionArgs {
661 args: args_values,
662 arg_fields,
663 number_rows: args.len(),
664 return_field: Field::new("f", Int64, true).into(),
665 })
666 }
667
668 fn test_case_sensitive_regexp_count_scalar() {
669 let values = ["", "aabca", "abcabc", "abcAbcab", "abcabcabc"];
670 let regex = "abc";
671 let expected: Vec<i64> = vec![0, 1, 2, 1, 3];
672
673 values.iter().enumerate().for_each(|(pos, &v)| {
674 let v_sv = ScalarValue::Utf8(Some(v.to_string()));
676 let regex_sv = ScalarValue::Utf8(Some(regex.to_string()));
677 let expected = expected.get(pos).cloned();
678 let re = regexp_count_with_scalar_values(&[v_sv, regex_sv]);
679 match re {
680 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
681 assert_eq!(v, expected, "regexp_count scalar test failed");
682 }
683 _ => panic!("Unexpected result"),
684 }
685
686 let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
688 let regex_sv = ScalarValue::LargeUtf8(Some(regex.to_string()));
689 let re = regexp_count_with_scalar_values(&[v_sv, regex_sv]);
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 let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
699 let regex_sv = ScalarValue::Utf8View(Some(regex.to_string()));
700 let re = regexp_count_with_scalar_values(&[v_sv, regex_sv]);
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 }
709
710 fn test_case_sensitive_regexp_count_scalar_start() {
711 let values = ["", "aabca", "abcabc", "abcAbcab", "abcabcabc"];
712 let regex = "abc";
713 let start = 2;
714 let expected: Vec<i64> = vec![0, 1, 1, 0, 2];
715
716 values.iter().enumerate().for_each(|(pos, &v)| {
717 let v_sv = ScalarValue::Utf8(Some(v.to_string()));
719 let regex_sv = ScalarValue::Utf8(Some(regex.to_string()));
720 let start_sv = ScalarValue::Int64(Some(start));
721 let expected = expected.get(pos).cloned();
722 let re = regexp_count_with_scalar_values(&[v_sv, regex_sv, start_sv.clone()]);
723 match re {
724 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
725 assert_eq!(v, expected, "regexp_count scalar test failed");
726 }
727 _ => panic!("Unexpected result"),
728 }
729
730 let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
732 let regex_sv = ScalarValue::LargeUtf8(Some(regex.to_string()));
733 let re = regexp_count_with_scalar_values(&[v_sv, regex_sv, start_sv.clone()]);
734 match re {
735 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
736 assert_eq!(v, expected, "regexp_count scalar test failed");
737 }
738 _ => panic!("Unexpected result"),
739 }
740
741 let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
743 let regex_sv = ScalarValue::Utf8View(Some(regex.to_string()));
744 let re = regexp_count_with_scalar_values(&[v_sv, regex_sv, start_sv.clone()]);
745 match re {
746 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
747 assert_eq!(v, expected, "regexp_count scalar test failed");
748 }
749 _ => panic!("Unexpected result"),
750 }
751 });
752 }
753
754 fn test_case_insensitive_regexp_count_scalar_flags() {
755 let values = ["", "aabca", "abcabc", "abcAbcab", "abcabcabc"];
756 let regex = "abc";
757 let start = 1;
758 let flags = "i";
759 let expected: Vec<i64> = vec![0, 1, 2, 2, 3];
760
761 values.iter().enumerate().for_each(|(pos, &v)| {
762 let v_sv = ScalarValue::Utf8(Some(v.to_string()));
764 let regex_sv = ScalarValue::Utf8(Some(regex.to_string()));
765 let start_sv = ScalarValue::Int64(Some(start));
766 let flags_sv = ScalarValue::Utf8(Some(flags.to_string()));
767 let expected = expected.get(pos).cloned();
768
769 let re = regexp_count_with_scalar_values(&[
770 v_sv,
771 regex_sv,
772 start_sv.clone(),
773 flags_sv.clone(),
774 ]);
775 match re {
776 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
777 assert_eq!(v, expected, "regexp_count scalar test failed");
778 }
779 _ => panic!("Unexpected result"),
780 }
781
782 let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
784 let regex_sv = ScalarValue::LargeUtf8(Some(regex.to_string()));
785 let flags_sv = ScalarValue::LargeUtf8(Some(flags.to_string()));
786
787 let re = regexp_count_with_scalar_values(&[
788 v_sv,
789 regex_sv,
790 start_sv.clone(),
791 flags_sv.clone(),
792 ]);
793 match re {
794 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
795 assert_eq!(v, expected, "regexp_count scalar test failed");
796 }
797 _ => panic!("Unexpected result"),
798 }
799
800 let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
802 let regex_sv = ScalarValue::Utf8View(Some(regex.to_string()));
803 let flags_sv = ScalarValue::Utf8View(Some(flags.to_string()));
804
805 let re = regexp_count_with_scalar_values(&[
806 v_sv,
807 regex_sv,
808 start_sv.clone(),
809 flags_sv.clone(),
810 ]);
811 match re {
812 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
813 assert_eq!(v, expected, "regexp_count scalar test failed");
814 }
815 _ => panic!("Unexpected result"),
816 }
817 });
818 }
819
820 fn test_case_sensitive_regexp_count_array<A>()
821 where
822 A: From<Vec<&'static str>> + Array + 'static,
823 {
824 let values = A::from(vec!["", "aabca", "abcabc", "abcAbcab", "abcabcAbc"]);
825 let regex = A::from(vec!["", "abc", "a", "bc", "ab"]);
826
827 let expected = Int64Array::from(vec![0, 1, 2, 2, 2]);
828
829 let re = regexp_count_func(&[Arc::new(values), Arc::new(regex)]).unwrap();
830 assert_eq!(re.as_ref(), &expected);
831 }
832
833 fn test_case_sensitive_regexp_count_array_start<A>()
834 where
835 A: From<Vec<&'static str>> + Array + 'static,
836 {
837 let values = A::from(vec!["", "aAbca", "abcabc", "abcAbcab", "abcabcAbc"]);
838 let regex = A::from(vec!["", "abc", "a", "bc", "ab"]);
839 let start = Int64Array::from(vec![1, 2, 3, 4, 5]);
840
841 let expected = Int64Array::from(vec![0, 0, 1, 1, 0]);
842
843 let re = regexp_count_func(&[Arc::new(values), Arc::new(regex), Arc::new(start)])
844 .unwrap();
845 assert_eq!(re.as_ref(), &expected);
846 }
847
848 fn test_case_insensitive_regexp_count_array_flags<A>()
849 where
850 A: From<Vec<&'static str>> + Array + 'static,
851 {
852 let values = A::from(vec!["", "aAbca", "abcabc", "abcAbcab", "abcabcAbc"]);
853 let regex = A::from(vec!["", "abc", "a", "bc", "ab"]);
854 let start = Int64Array::from(vec![1]);
855 let flags = A::from(vec!["", "i", "", "", "i"]);
856
857 let expected = Int64Array::from(vec![0, 1, 2, 2, 3]);
858
859 let re = regexp_count_func(&[
860 Arc::new(values),
861 Arc::new(regex),
862 Arc::new(start),
863 Arc::new(flags),
864 ])
865 .unwrap();
866 assert_eq!(re.as_ref(), &expected);
867 }
868
869 fn test_case_sensitive_regexp_count_start_scalar_complex() {
870 let values = ["", "aabca", "abcabc", "abcAbcab", "abcabcabc"];
871 let regex = ["", "abc", "a", "bc", "ab"];
872 let start = 5;
873 let flags = ["", "i", "", "", "i"];
874 let expected: Vec<i64> = vec![0, 0, 0, 1, 1];
875
876 values.iter().enumerate().for_each(|(pos, &v)| {
877 let v_sv = ScalarValue::Utf8(Some(v.to_string()));
879 let regex_sv = ScalarValue::Utf8(regex.get(pos).map(|s| s.to_string()));
880 let start_sv = ScalarValue::Int64(Some(start));
881 let flags_sv = ScalarValue::Utf8(flags.get(pos).map(|f| f.to_string()));
882 let expected = expected.get(pos).cloned();
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 let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
898 let regex_sv = ScalarValue::LargeUtf8(regex.get(pos).map(|s| s.to_string()));
899 let flags_sv = ScalarValue::LargeUtf8(flags.get(pos).map(|f| f.to_string()));
900 let re = regexp_count_with_scalar_values(&[
901 v_sv,
902 regex_sv,
903 start_sv.clone(),
904 flags_sv.clone(),
905 ]);
906 match re {
907 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
908 assert_eq!(v, expected, "regexp_count scalar test failed");
909 }
910 _ => panic!("Unexpected result"),
911 }
912
913 let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
915 let regex_sv = ScalarValue::Utf8View(regex.get(pos).map(|s| s.to_string()));
916 let flags_sv = ScalarValue::Utf8View(flags.get(pos).map(|f| f.to_string()));
917 let re = regexp_count_with_scalar_values(&[
918 v_sv,
919 regex_sv,
920 start_sv.clone(),
921 flags_sv.clone(),
922 ]);
923 match re {
924 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
925 assert_eq!(v, expected, "regexp_count scalar test failed");
926 }
927 _ => panic!("Unexpected result"),
928 }
929 });
930 }
931
932 fn test_case_sensitive_regexp_count_array_complex<A>()
933 where
934 A: From<Vec<&'static str>> + Array + 'static,
935 {
936 let values = A::from(vec!["", "aAbca", "abcabc", "abcAbcab", "abcabcAbc"]);
937 let regex = A::from(vec!["", "abc", "a", "bc", "ab"]);
938 let start = Int64Array::from(vec![1, 2, 3, 4, 5]);
939 let flags = A::from(vec!["", "i", "", "", "i"]);
940
941 let expected = Int64Array::from(vec![0, 1, 1, 1, 1]);
942
943 let re = regexp_count_func(&[
944 Arc::new(values),
945 Arc::new(regex),
946 Arc::new(start),
947 Arc::new(flags),
948 ])
949 .unwrap();
950 assert_eq!(re.as_ref(), &expected);
951 }
952
953 fn test_case_regexp_count_cache_check<A>()
954 where
955 A: From<Vec<&'static str>> + Array + 'static,
956 {
957 let values = A::from(vec!["aaa", "Aaa", "aaa"]);
958 let regex = A::from(vec!["aaa", "aaa", "aaa"]);
959 let start = Int64Array::from(vec![1, 1, 1]);
960 let flags = A::from(vec!["", "i", ""]);
961
962 let expected = Int64Array::from(vec![1, 1, 1]);
963
964 let re = regexp_count_func(&[
965 Arc::new(values),
966 Arc::new(regex),
967 Arc::new(start),
968 Arc::new(flags),
969 ])
970 .unwrap();
971 assert_eq!(re.as_ref(), &expected);
972 }
973}