1use 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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}