1use arrow::array::{
19 Array, ArrayRef, AsArray, Datum, Int64Array, Int64Builder, StringArrayType,
20};
21use arrow::buffer::NullBuffer;
22use arrow::datatypes::{DataType, Int64Type};
23use arrow::datatypes::{
24 DataType::Int64, DataType::LargeUtf8, DataType::Utf8, DataType::Utf8View,
25};
26use arrow::error::ArrowError;
27use datafusion_common::{Result, ScalarValue, exec_err, internal_err};
28use datafusion_expr::{
29 ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
30 TypeSignature::Exact, TypeSignature::Uniform, Volatility,
31};
32use datafusion_macros::user_doc;
33use regex::Regex;
34use std::collections::HashMap;
35use std::collections::hash_map::Entry;
36use std::sync::Arc;
37
38use crate::regex::{compile_regex, start_to_byte_offset};
39
40#[user_doc(
41 doc_section(label = "Regular Expression Functions"),
42 description = "Returns the position in a string where the specified occurrence of a POSIX regular expression is located.",
43 syntax_example = "regexp_instr(str, regexp[, start[, N[, flags[, subexpr]]]])",
44 sql_example = r#"```sql
45> SELECT regexp_instr('ABCDEF', 'C(.)(..)');
46+---------------------------------------------------------------+
47| regexp_instr(Utf8("ABCDEF"),Utf8("C(.)(..)")) |
48+---------------------------------------------------------------+
49| 3 |
50+---------------------------------------------------------------+
51```"#,
52 standard_argument(name = "str", prefix = "String"),
53 standard_argument(name = "regexp", prefix = "Regular"),
54 argument(
55 name = "start",
56 description = "Optional start position (the first position is 1) to search for the regular expression. Can be a constant, column, or function. Defaults to 1"
57 ),
58 argument(
59 name = "N",
60 description = "Optional The N-th occurrence of pattern to find. Defaults to 1 (first match). Can be a constant, column, or function."
61 ),
62 argument(
63 name = "flags",
64 description = r#"Optional regular expression flags that control the behavior of the regular expression. Refer to the flags reference above for supported flags."#
65 ),
66 argument(
67 name = "subexpr",
68 description = "Optional Specifies which capture group (subexpression) to return the position for. Defaults to 0, which returns the position of the entire match."
69 )
70)]
71#[derive(Debug, PartialEq, Eq, Hash)]
72pub struct RegexpInstrFunc {
73 signature: Signature,
74}
75
76impl Default for RegexpInstrFunc {
77 fn default() -> Self {
78 Self::new()
79 }
80}
81
82impl RegexpInstrFunc {
83 pub fn new() -> Self {
84 Self {
85 signature: Signature::one_of(
86 vec![
87 Uniform(2, vec![Utf8View, LargeUtf8, Utf8]),
88 Exact(vec![Utf8View, Utf8View, Int64]),
89 Exact(vec![LargeUtf8, LargeUtf8, Int64]),
90 Exact(vec![Utf8, Utf8, Int64]),
91 Exact(vec![Utf8View, Utf8View, Int64, Int64]),
92 Exact(vec![LargeUtf8, LargeUtf8, Int64, Int64]),
93 Exact(vec![Utf8, Utf8, Int64, Int64]),
94 Exact(vec![Utf8View, Utf8View, Int64, Int64, Utf8View]),
95 Exact(vec![LargeUtf8, LargeUtf8, Int64, Int64, LargeUtf8]),
96 Exact(vec![Utf8, Utf8, Int64, Int64, Utf8]),
97 Exact(vec![Utf8View, Utf8View, Int64, Int64, Utf8View, Int64]),
98 Exact(vec![LargeUtf8, LargeUtf8, Int64, Int64, LargeUtf8, Int64]),
99 Exact(vec![Utf8, Utf8, Int64, Int64, Utf8, Int64]),
100 ],
101 Volatility::Immutable,
102 ),
103 }
104 }
105}
106
107impl ScalarUDFImpl for RegexpInstrFunc {
108 fn name(&self) -> &str {
109 "regexp_instr"
110 }
111
112 fn signature(&self) -> &Signature {
113 &self.signature
114 }
115
116 fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
117 Ok(Int64)
118 }
119
120 fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
121 let args = &args.args;
122
123 let len = args
124 .iter()
125 .fold(Option::<usize>::None, |acc, arg| match arg {
126 ColumnarValue::Scalar(_) => acc,
127 ColumnarValue::Array(a) => Some(a.len()),
128 });
129
130 let is_scalar = len.is_none();
131 let inferred_length = len.unwrap_or(1);
132 let args = args
133 .iter()
134 .map(|arg| arg.to_array(inferred_length))
135 .collect::<Result<Vec<_>>>()?;
136
137 let result = regexp_instr_func(&args);
138 if is_scalar {
139 let result = result.and_then(|arr| ScalarValue::try_from_array(&arr, 0));
141 result.map(ColumnarValue::Scalar)
142 } else {
143 result.map(ColumnarValue::Array)
144 }
145 }
146
147 fn documentation(&self) -> Option<&Documentation> {
148 self.doc()
149 }
150}
151
152pub fn regexp_instr_func(args: &[ArrayRef]) -> Result<ArrayRef> {
153 let args_len = args.len();
154 if !(2..=6).contains(&args_len) {
155 return exec_err!(
156 "regexp_instr was called with {args_len} arguments. It requires at least 2 and at most 6."
157 );
158 }
159
160 let values = &args[0];
161 match values.data_type() {
162 Utf8 | LargeUtf8 | Utf8View => (),
163 other => {
164 return internal_err!(
165 "Unsupported data type {other:?} for function regexp_instr"
166 );
167 }
168 }
169
170 regexp_instr(
171 values,
172 &args[1],
173 if args_len > 2 { Some(&args[2]) } else { None },
174 if args_len > 3 { Some(&args[3]) } else { None },
175 if args_len > 4 { Some(&args[4]) } else { None },
176 if args_len > 5 { Some(&args[5]) } else { None },
177 )
178 .map_err(|e| e.into())
179}
180
181fn regexp_instr(
200 values: &dyn Array,
201 regex_array: &dyn Datum,
202 start_array: Option<&dyn Datum>,
203 nth_array: Option<&dyn Datum>,
204 flags_array: Option<&dyn Datum>,
205 subexpr_array: Option<&dyn Datum>,
206) -> Result<ArrayRef, ArrowError> {
207 let (regex_array, _) = regex_array.get();
208 let start_array = start_array.map(|start| {
209 let (start, _) = start.get();
210 start
211 });
212 let nth_array = nth_array.map(|nth| {
213 let (nth, _) = nth.get();
214 nth
215 });
216 let flags_array = flags_array.map(|flags| {
217 let (flags, _) = flags.get();
218 flags
219 });
220 let subexpr_array = subexpr_array.map(|subexpr| {
221 let (subexpr, _) = subexpr.get();
222 subexpr
223 });
224
225 match (values.data_type(), regex_array.data_type(), flags_array) {
226 (Utf8, Utf8, None) => regexp_instr_inner(
227 &values.as_string::<i32>(),
228 ®ex_array.as_string::<i32>(),
229 start_array.map(|start| start.as_primitive::<Int64Type>()),
230 nth_array.map(|nth| nth.as_primitive::<Int64Type>()),
231 None,
232 subexpr_array.map(|subexpr| subexpr.as_primitive::<Int64Type>()),
233 ),
234 (Utf8, Utf8, Some(flags_array)) if *flags_array.data_type() == Utf8 => regexp_instr_inner(
235 &values.as_string::<i32>(),
236 ®ex_array.as_string::<i32>(),
237 start_array.map(|start| start.as_primitive::<Int64Type>()),
238 nth_array.map(|nth| nth.as_primitive::<Int64Type>()),
239 Some(&flags_array.as_string::<i32>()),
240 subexpr_array.map(|subexpr| subexpr.as_primitive::<Int64Type>()),
241 ),
242 (LargeUtf8, LargeUtf8, None) => regexp_instr_inner(
243 &values.as_string::<i64>(),
244 ®ex_array.as_string::<i64>(),
245 start_array.map(|start| start.as_primitive::<Int64Type>()),
246 nth_array.map(|nth| nth.as_primitive::<Int64Type>()),
247 None,
248 subexpr_array.map(|subexpr| subexpr.as_primitive::<Int64Type>()),
249 ),
250 (LargeUtf8, LargeUtf8, Some(flags_array)) if *flags_array.data_type() == LargeUtf8 => regexp_instr_inner(
251 &values.as_string::<i64>(),
252 ®ex_array.as_string::<i64>(),
253 start_array.map(|start| start.as_primitive::<Int64Type>()),
254 nth_array.map(|nth| nth.as_primitive::<Int64Type>()),
255 Some(&flags_array.as_string::<i64>()),
256 subexpr_array.map(|subexpr| subexpr.as_primitive::<Int64Type>()),
257 ),
258 (Utf8View, Utf8View, None) => regexp_instr_inner(
259 &values.as_string_view(),
260 ®ex_array.as_string_view(),
261 start_array.map(|start| start.as_primitive::<Int64Type>()),
262 nth_array.map(|nth| nth.as_primitive::<Int64Type>()),
263 None,
264 subexpr_array.map(|subexpr| subexpr.as_primitive::<Int64Type>()),
265 ),
266 (Utf8View, Utf8View, Some(flags_array)) if *flags_array.data_type() == Utf8View => regexp_instr_inner(
267 &values.as_string_view(),
268 ®ex_array.as_string_view(),
269 start_array.map(|start| start.as_primitive::<Int64Type>()),
270 nth_array.map(|nth| nth.as_primitive::<Int64Type>()),
271 Some(&flags_array.as_string_view()),
272 subexpr_array.map(|subexpr| subexpr.as_primitive::<Int64Type>()),
273 ),
274 _ => Err(ArrowError::ComputeError(
275 "regexp_instr() 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(),
276 )),
277 }
278}
279
280fn regexp_instr_inner<'a, S>(
281 values: &S,
282 regex_array: &S,
283 start_array: Option<&Int64Array>,
284 nth_array: Option<&Int64Array>,
285 flags_array: Option<&S>,
286 subexp_array: Option<&Int64Array>,
287) -> Result<ArrayRef, ArrowError>
288where
289 S: StringArrayType<'a>,
290{
291 let len = values.len();
292 let mut regex_cache = RegexCache::default();
293 let mut result = Int64Builder::with_capacity(len);
294
295 let nulls = NullBuffer::union_many([
297 values.nulls(),
298 regex_array.nulls(),
299 start_array.and_then(|array| array.nulls()),
300 nth_array.and_then(|array| array.nulls()),
301 flags_array.and_then(|array| array.nulls()),
302 subexp_array.and_then(|array| array.nulls()),
303 ]);
304
305 for i in 0..len {
306 if nulls.as_ref().is_some_and(|nulls| nulls.is_null(i)) {
307 result.append_null();
308 continue;
309 }
310
311 let value = values.value(i);
312 let regex = regex_array.value(i);
313 let flags = flags_array.map(|array| array.value(i));
314 let pattern = regex_cache.get_or_compile(regex, flags)?;
315
316 let start = start_array.map_or(1, |array| array.value(i));
318 let nth = nth_array.map_or(1, |array| array.value(i));
319 let subexp = subexp_array.map_or(0, |array| array.value(i));
320
321 result.append_value(get_index(value, pattern, start, nth, subexp)?);
322 }
323
324 Ok(Arc::new(result.finish()))
325}
326
327#[derive(Default)]
334struct RegexCache<'a> {
335 compiled: Vec<Regex>,
336 indices: HashMap<(&'a str, Option<&'a str>), usize>,
337 last: Option<((&'a str, Option<&'a str>), usize)>,
338}
339
340impl<'a> RegexCache<'a> {
341 fn get_or_compile(
342 &mut self,
343 regex: &'a str,
344 flags: Option<&'a str>,
345 ) -> Result<&Regex, ArrowError> {
346 let key = (regex, flags);
347 let index = match self.last {
348 Some((last_key, index)) if last_key == key => index,
349 _ => {
350 let index = match self.indices.entry(key) {
351 Entry::Occupied(entry) => *entry.get(),
352 Entry::Vacant(entry) => {
353 self.compiled.push(compile_regex(regex, flags)?);
354 *entry.insert(self.compiled.len() - 1)
355 }
356 };
357 self.last = Some((key, index));
358 index
359 }
360 };
361 Ok(&self.compiled[index])
362 }
363}
364
365fn get_index(
370 value: &str,
371 pattern: &Regex,
372 start: i64,
373 n: i64,
374 subexpr: i64,
375) -> Result<i64, ArrowError> {
376 if start < 1 {
377 return Err(ArrowError::ComputeError(
378 "regexp_instr() requires start to be 1-based".to_string(),
379 ));
380 }
381
382 if n < 1 {
383 return Err(ArrowError::ComputeError(
384 "N must be 1 or greater".to_string(),
385 ));
386 }
387
388 let Some(byte_start_offset) = start_to_byte_offset(value, start) else {
389 return Ok(0);
390 };
391 let search_slice = &value[byte_start_offset..];
392
393 let match_start = if subexpr > 0 {
395 pattern
396 .captures(search_slice)
397 .and_then(|captures| captures.get(subexpr as usize))
398 .map(|matched| matched.start())
399 } else {
400 pattern
402 .find_iter(search_slice)
403 .nth((n - 1) as usize)
404 .map(|matched| matched.start())
405 };
406
407 Ok(match_start.map_or(0, |offset| {
410 value[..byte_start_offset + offset].chars().count() as i64 + 1
411 }))
412}
413
414#[cfg(test)]
415mod tests {
416 use super::*;
417 use arrow::array::{GenericStringArray, StringViewArray};
418 use arrow::datatypes::Field;
419 use datafusion_common::config::ConfigOptions;
420 use itertools::izip;
421 #[test]
422 fn test_regexp_instr() {
423 test_case_sensitive_regexp_instr_nulls();
424 test_case_sensitive_regexp_instr_scalar();
425 test_case_sensitive_regexp_instr_scalar_start();
426 test_case_sensitive_regexp_instr_scalar_nth();
427 test_case_sensitive_regexp_instr_scalar_subexp();
428
429 test_case_sensitive_regexp_instr_array::<GenericStringArray<i32>>();
430 test_case_sensitive_regexp_instr_array::<GenericStringArray<i64>>();
431 test_case_sensitive_regexp_instr_array::<StringViewArray>();
432
433 test_case_sensitive_regexp_instr_array_start::<GenericStringArray<i32>>();
434 test_case_sensitive_regexp_instr_array_start::<GenericStringArray<i64>>();
435 test_case_sensitive_regexp_instr_array_start::<StringViewArray>();
436
437 test_case_sensitive_regexp_instr_array_nth::<GenericStringArray<i32>>();
438 test_case_sensitive_regexp_instr_array_nth::<GenericStringArray<i64>>();
439 test_case_sensitive_regexp_instr_array_nth::<StringViewArray>();
440
441 test_case_sensitive_regexp_instr_empty_pattern::<GenericStringArray<i32>>();
442 test_case_sensitive_regexp_instr_empty_pattern::<GenericStringArray<i64>>();
443 test_case_sensitive_regexp_instr_empty_pattern::<StringViewArray>();
444
445 test_case_sensitive_regexp_instr_zero_width_pattern::<GenericStringArray<i32>>();
446 test_case_sensitive_regexp_instr_zero_width_pattern::<GenericStringArray<i64>>();
447 test_case_sensitive_regexp_instr_zero_width_pattern::<StringViewArray>();
448
449 test_regexp_instr_null_scalar_args();
450
451 test_regexp_instr_null_array_rows::<GenericStringArray<i32>>();
452 test_regexp_instr_null_array_rows::<GenericStringArray<i64>>();
453 test_regexp_instr_null_array_rows::<StringViewArray>();
454 }
455
456 fn regexp_instr_with_scalar_values(args: &[ScalarValue]) -> Result<ColumnarValue> {
457 let args_values: Vec<ColumnarValue> = args
458 .iter()
459 .map(|sv| ColumnarValue::Scalar(sv.clone()))
460 .collect();
461
462 let arg_fields = args
463 .iter()
464 .enumerate()
465 .map(|(idx, a)| {
466 Arc::new(Field::new(format!("arg_{idx}"), a.data_type(), true))
467 })
468 .collect::<Vec<_>>();
469
470 RegexpInstrFunc::new().invoke_with_args(ScalarFunctionArgs {
471 args: args_values,
472 arg_fields,
473 number_rows: args.len(),
474 return_field: Arc::new(Field::new("f", Int64, true)),
475 config_options: Arc::new(ConfigOptions::default()),
476 })
477 }
478
479 fn test_case_sensitive_regexp_instr_nulls() {
480 let v = "";
481 let r = "";
482 let expected = 1;
483 let regex_sv = ScalarValue::Utf8(Some(r.to_string()));
484 let re = regexp_instr_with_scalar_values(&[v.to_string().into(), regex_sv]);
485 match re {
487 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
488 assert_eq!(v, Some(expected), "regexp_instr scalar test failed");
489 }
490 _ => panic!("Unexpected result"),
491 }
492
493 for (value, regex) in [
494 (
495 ScalarValue::Utf8(None),
496 ScalarValue::Utf8(Some(String::new())),
497 ),
498 (
499 ScalarValue::LargeUtf8(None),
500 ScalarValue::LargeUtf8(Some(String::new())),
501 ),
502 (
503 ScalarValue::Utf8View(None),
504 ScalarValue::Utf8View(Some(String::new())),
505 ),
506 ] {
507 let re = regexp_instr_with_scalar_values(&[value, regex]);
508 match re {
509 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
510 assert_eq!(v, None, "regexp_instr NULL scalar test failed");
511 }
512 _ => panic!("Unexpected result"),
513 }
514 }
515 }
516 fn test_case_sensitive_regexp_instr_scalar() {
517 let values = [
518 "hello world",
519 "abcdefg",
520 "xyz123xyz",
521 "no match here",
522 "abc",
523 "ДатаФусион数据融合📊🔥",
524 ];
525 let regex = ["o", "d", "123", "z", "gg", "📊"];
526
527 let expected: Vec<i64> = vec![5, 4, 4, 0, 0, 15];
528
529 izip!(values.iter(), regex.iter())
530 .enumerate()
531 .for_each(|(pos, (&v, &r))| {
532 let v_sv = ScalarValue::Utf8(Some(v.to_string()));
534 let regex_sv = ScalarValue::Utf8(Some(r.to_string()));
535 let expected = expected.get(pos).cloned();
536 let re = regexp_instr_with_scalar_values(&[v_sv, regex_sv]);
537 match re {
539 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
540 assert_eq!(v, expected, "regexp_instr scalar test failed");
541 }
542 _ => panic!("Unexpected result"),
543 }
544
545 let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
547 let regex_sv = ScalarValue::LargeUtf8(Some(r.to_string()));
548 let re = regexp_instr_with_scalar_values(&[v_sv, regex_sv]);
549 match re {
550 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
551 assert_eq!(v, expected, "regexp_instr scalar test failed");
552 }
553 _ => panic!("Unexpected result"),
554 }
555
556 let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
558 let regex_sv = ScalarValue::Utf8View(Some(r.to_string()));
559 let re = regexp_instr_with_scalar_values(&[v_sv, regex_sv]);
560 match re {
561 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
562 assert_eq!(v, expected, "regexp_instr scalar test failed");
563 }
564 _ => panic!("Unexpected result"),
565 }
566 });
567 }
568
569 fn test_case_sensitive_regexp_instr_scalar_start() {
570 let values = ["abcabcabc", "abcabcabc", ""];
571 let regex = ["abc", "abc", "gg"];
572 let start = [4, 5, 5];
573 let expected: Vec<i64> = vec![4, 7, 0];
574
575 izip!(values.iter(), regex.iter(), start.iter())
576 .enumerate()
577 .for_each(|(pos, (&v, &r, &s))| {
578 let v_sv = ScalarValue::Utf8(Some(v.to_string()));
580 let regex_sv = ScalarValue::Utf8(Some(r.to_string()));
581 let start_sv = ScalarValue::Int64(Some(s));
582 let expected = expected.get(pos).cloned();
583 let re =
584 regexp_instr_with_scalar_values(&[v_sv, regex_sv, start_sv.clone()]);
585 match re {
586 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
587 assert_eq!(v, expected, "regexp_instr scalar test failed");
588 }
589 _ => panic!("Unexpected result"),
590 }
591
592 let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
594 let regex_sv = ScalarValue::LargeUtf8(Some(r.to_string()));
595 let start_sv = ScalarValue::Int64(Some(s));
596 let re =
597 regexp_instr_with_scalar_values(&[v_sv, regex_sv, start_sv.clone()]);
598 match re {
599 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
600 assert_eq!(v, expected, "regexp_instr scalar test failed");
601 }
602 _ => panic!("Unexpected result"),
603 }
604
605 let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
607 let regex_sv = ScalarValue::Utf8View(Some(r.to_string()));
608 let start_sv = ScalarValue::Int64(Some(s));
609 let re =
610 regexp_instr_with_scalar_values(&[v_sv, regex_sv, start_sv.clone()]);
611 match re {
612 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
613 assert_eq!(v, expected, "regexp_instr scalar test failed");
614 }
615 _ => panic!("Unexpected result"),
616 }
617 });
618 }
619
620 fn test_case_sensitive_regexp_instr_scalar_nth() {
621 let values = ["abcabcabc", "abcabcabc", "abcabcabc", "abcabcabc"];
622 let regex = ["abc", "abc", "abc", "abc"];
623 let start = [1, 1, 1, 1];
624 let nth = [1, 2, 3, 4];
625 let expected: Vec<i64> = vec![1, 4, 7, 0];
626
627 izip!(values.iter(), regex.iter(), start.iter(), nth.iter())
628 .enumerate()
629 .for_each(|(pos, (&v, &r, &s, &n))| {
630 let v_sv = ScalarValue::Utf8(Some(v.to_string()));
632 let regex_sv = ScalarValue::Utf8(Some(r.to_string()));
633 let start_sv = ScalarValue::Int64(Some(s));
634 let nth_sv = ScalarValue::Int64(Some(n));
635 let expected = expected.get(pos).cloned();
636 let re = regexp_instr_with_scalar_values(&[
637 v_sv,
638 regex_sv,
639 start_sv.clone(),
640 nth_sv.clone(),
641 ]);
642 match re {
643 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
644 assert_eq!(v, expected, "regexp_instr scalar test failed");
645 }
646 _ => panic!("Unexpected result"),
647 }
648
649 let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
651 let regex_sv = ScalarValue::LargeUtf8(Some(r.to_string()));
652 let start_sv = ScalarValue::Int64(Some(s));
653 let nth_sv = ScalarValue::Int64(Some(n));
654 let re = regexp_instr_with_scalar_values(&[
655 v_sv,
656 regex_sv,
657 start_sv.clone(),
658 nth_sv.clone(),
659 ]);
660 match re {
661 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
662 assert_eq!(v, expected, "regexp_instr scalar test failed");
663 }
664 _ => panic!("Unexpected result"),
665 }
666
667 let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
669 let regex_sv = ScalarValue::Utf8View(Some(r.to_string()));
670 let start_sv = ScalarValue::Int64(Some(s));
671 let nth_sv = ScalarValue::Int64(Some(n));
672 let re = regexp_instr_with_scalar_values(&[
673 v_sv,
674 regex_sv,
675 start_sv.clone(),
676 nth_sv.clone(),
677 ]);
678 match re {
679 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
680 assert_eq!(v, expected, "regexp_instr scalar test failed");
681 }
682 _ => panic!("Unexpected result"),
683 }
684 });
685 }
686
687 fn test_case_sensitive_regexp_instr_scalar_subexp() {
688 let values = ["12 abc def ghi 34"];
689 let regex = ["(abc) (def) (ghi)"];
690 let start = [1];
691 let nth = [1];
692 let flags = ["i"];
693 let subexps = [2];
694 let expected: Vec<i64> = vec![8];
695
696 izip!(
697 values.iter(),
698 regex.iter(),
699 start.iter(),
700 nth.iter(),
701 flags.iter(),
702 subexps.iter()
703 )
704 .enumerate()
705 .for_each(|(pos, (&v, &r, &s, &n, &flag, &subexp))| {
706 let v_sv = ScalarValue::Utf8(Some(v.to_string()));
708 let regex_sv = ScalarValue::Utf8(Some(r.to_string()));
709 let start_sv = ScalarValue::Int64(Some(s));
710 let nth_sv = ScalarValue::Int64(Some(n));
711 let flags_sv = ScalarValue::Utf8(Some(flag.to_string()));
712 let subexp_sv = ScalarValue::Int64(Some(subexp));
713 let expected = expected.get(pos).cloned();
714 let re = regexp_instr_with_scalar_values(&[
715 v_sv,
716 regex_sv,
717 start_sv.clone(),
718 nth_sv.clone(),
719 flags_sv,
720 subexp_sv.clone(),
721 ]);
722 match re {
723 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
724 assert_eq!(v, expected, "regexp_instr scalar test failed");
725 }
726 _ => panic!("Unexpected result"),
727 }
728
729 let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
731 let regex_sv = ScalarValue::LargeUtf8(Some(r.to_string()));
732 let start_sv = ScalarValue::Int64(Some(s));
733 let nth_sv = ScalarValue::Int64(Some(n));
734 let flags_sv = ScalarValue::LargeUtf8(Some(flag.to_string()));
735 let subexp_sv = ScalarValue::Int64(Some(subexp));
736 let re = regexp_instr_with_scalar_values(&[
737 v_sv,
738 regex_sv,
739 start_sv.clone(),
740 nth_sv.clone(),
741 flags_sv,
742 subexp_sv.clone(),
743 ]);
744 match re {
745 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
746 assert_eq!(v, expected, "regexp_instr scalar test failed");
747 }
748 _ => panic!("Unexpected result"),
749 }
750
751 let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
753 let regex_sv = ScalarValue::Utf8View(Some(r.to_string()));
754 let start_sv = ScalarValue::Int64(Some(s));
755 let nth_sv = ScalarValue::Int64(Some(n));
756 let flags_sv = ScalarValue::Utf8View(Some(flag.to_string()));
757 let subexp_sv = ScalarValue::Int64(Some(subexp));
758 let re = regexp_instr_with_scalar_values(&[
759 v_sv,
760 regex_sv,
761 start_sv.clone(),
762 nth_sv.clone(),
763 flags_sv,
764 subexp_sv.clone(),
765 ]);
766 match re {
767 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
768 assert_eq!(v, expected, "regexp_instr scalar test failed");
769 }
770 _ => panic!("Unexpected result"),
771 }
772 });
773 }
774
775 fn test_regexp_instr_null_scalar_args() {
776 let cases: Vec<Vec<ScalarValue>> = vec![
778 vec![
780 ScalarValue::Utf8(Some("abc".to_string())),
781 ScalarValue::Utf8(Some("b".to_string())),
782 ScalarValue::Int64(None),
783 ],
784 vec![
786 ScalarValue::Utf8(Some("abc".to_string())),
787 ScalarValue::Utf8(Some("b".to_string())),
788 ScalarValue::Int64(Some(1)),
789 ScalarValue::Int64(None),
790 ],
791 vec![
793 ScalarValue::Utf8(Some("abc".to_string())),
794 ScalarValue::Utf8(Some("b".to_string())),
795 ScalarValue::Int64(Some(1)),
796 ScalarValue::Int64(Some(1)),
797 ScalarValue::Utf8(None),
798 ],
799 vec![
801 ScalarValue::Utf8(Some("abc".to_string())),
802 ScalarValue::Utf8(Some("(b)".to_string())),
803 ScalarValue::Int64(Some(1)),
804 ScalarValue::Int64(Some(1)),
805 ScalarValue::Utf8(Some("i".to_string())),
806 ScalarValue::Int64(None),
807 ],
808 ];
809
810 for args in cases {
811 let re = regexp_instr_with_scalar_values(&args);
812 match re {
813 Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
814 assert_eq!(v, None, "regexp_instr null scalar test failed");
815 }
816 _ => panic!("Unexpected result"),
817 }
818 }
819 }
820
821 fn test_regexp_instr_null_array_rows<A>()
822 where
823 A: From<Vec<Option<&'static str>>> + Array + 'static,
824 {
825 let values = A::from(vec![
826 None,
827 Some("abc"),
828 Some("abc"),
829 Some("abc"),
830 Some("abc"),
831 Some("abc"),
832 Some("abc"),
833 ]);
834 let regex = A::from(vec![
835 Some("b"),
836 None,
837 Some("b"),
838 Some("b"),
839 Some("b"),
840 Some("(b)"),
841 Some("b"),
842 ]);
843 let start = Int64Array::from(vec![
844 Some(1),
845 Some(1),
846 None,
847 Some(1),
848 Some(1),
849 Some(1),
850 Some(1),
851 ]);
852 let nth = Int64Array::from(vec![
853 Some(1),
854 Some(1),
855 Some(1),
856 None,
857 Some(1),
858 Some(1),
859 Some(1),
860 ]);
861 let flags = A::from(vec![
862 Some(""),
863 Some(""),
864 Some(""),
865 Some(""),
866 None,
867 Some("i"),
868 Some(""),
869 ]);
870 let subexp = Int64Array::from(vec![
871 Some(0),
872 Some(0),
873 Some(0),
874 Some(0),
875 Some(0),
876 None,
877 Some(0),
878 ]);
879
880 let expected =
881 Int64Array::from(vec![None, None, None, None, None, None, Some(2)]);
882
883 let re = regexp_instr_func(&[
884 Arc::new(values),
885 Arc::new(regex),
886 Arc::new(start),
887 Arc::new(nth),
888 Arc::new(flags),
889 Arc::new(subexp),
890 ])
891 .unwrap();
892 assert_eq!(re.as_ref(), &expected);
893 }
894
895 fn test_case_sensitive_regexp_instr_array<A>()
896 where
897 A: From<Vec<&'static str>> + Array + 'static,
898 {
899 let values = A::from(vec![
900 "hello world",
901 "abcdefg",
902 "xyz123xyz",
903 "no match here",
904 "",
905 ]);
906 let regex = A::from(vec!["o", "d", "123", "z", "gg"]);
907
908 let expected = Int64Array::from(vec![5, 4, 4, 0, 0]);
909 let re = regexp_instr_func(&[Arc::new(values), Arc::new(regex)]).unwrap();
910 assert_eq!(re.as_ref(), &expected);
911 }
912
913 fn test_case_sensitive_regexp_instr_array_start<A>()
914 where
915 A: From<Vec<&'static str>> + Array + 'static,
916 {
917 let values = A::from(vec!["abcabcabc", "abcabcabc", ""]);
918 let regex = A::from(vec!["abc", "abc", "gg"]);
919 let start = Int64Array::from(vec![4, 5, 5]);
920 let expected = Int64Array::from(vec![4, 7, 0]);
921
922 let re = regexp_instr_func(&[Arc::new(values), Arc::new(regex), Arc::new(start)])
923 .unwrap();
924 assert_eq!(re.as_ref(), &expected);
925 }
926
927 fn test_case_sensitive_regexp_instr_array_nth<A>()
928 where
929 A: From<Vec<&'static str>> + Array + 'static,
930 {
931 let values = A::from(vec!["abcabcabc", "abcabcabc", "abcabcabc", "abcabcabc"]);
932 let regex = A::from(vec!["abc", "abc", "abc", "abc"]);
933 let start = Int64Array::from(vec![1, 1, 1, 1]);
934 let nth = Int64Array::from(vec![1, 2, 3, 4]);
935 let expected = Int64Array::from(vec![1, 4, 7, 0]);
936
937 let re = regexp_instr_func(&[
938 Arc::new(values),
939 Arc::new(regex),
940 Arc::new(start),
941 Arc::new(nth),
942 ])
943 .unwrap();
944 assert_eq!(re.as_ref(), &expected);
945 }
946
947 fn test_case_sensitive_regexp_instr_empty_pattern<A>()
948 where
949 A: From<Vec<&'static str>> + Array + 'static,
950 {
951 let values = A::from(vec!["abc", "", "abc", "abc", "😀"]);
952 let regex = A::from(vec!["", "", "", "", ""]);
953 let start = Int64Array::from(vec![1, 1, 4, 5, 1]);
954 let nth = Int64Array::from(vec![1, 1, 1, 1, 2]);
955 let expected = Int64Array::from(vec![1, 1, 4, 0, 2]);
956
957 let re = regexp_instr_func(&[
958 Arc::new(values),
959 Arc::new(regex),
960 Arc::new(start),
961 Arc::new(nth),
962 ])
963 .unwrap();
964 assert_eq!(re.as_ref(), &expected);
965 }
966
967 fn test_case_sensitive_regexp_instr_zero_width_pattern<A>()
968 where
969 A: From<Vec<&'static str>> + Array + 'static,
970 {
971 let values = A::from(vec!["abc"]);
972 let regex = A::from(vec!["x*"]);
973 let start = Int64Array::from(vec![4]);
974 let expected = Int64Array::from(vec![4]);
975
976 let re = regexp_instr_func(&[Arc::new(values), Arc::new(regex), Arc::new(start)])
977 .unwrap();
978 assert_eq!(re.as_ref(), &expected);
979 }
980}