Skip to main content

datafusion_functions/regex/
regexpinstr.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 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            // If all inputs are scalar, keeps output as scalar
140            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
181/// `arrow-rs` style implementation of `regexp_instr` function.
182/// This function `regexp_instr` is responsible for returning the index of a regular expression pattern
183/// within a string array. It supports optional start positions and flags for case insensitivity.
184///
185/// The function accepts a variable number of arguments:
186/// - `values`: The array of strings to search within.
187/// - `regex_array`: The array of regular expression patterns to search for.
188/// - `start_array` (optional): The array of start positions for the search.
189/// - `nth_array` (optional): The array of start nth for the search.
190/// - `endoption_array` (optional): The array of endoption positions for the search.
191/// - `flags_array` (optional): The array of flags to modify the search behavior (e.g., case insensitivity).
192/// - `subexpr_array` (optional): The array of subexpr positions for the search.
193///
194/// The function handles different combinations of scalar and array inputs for the regex patterns, start positions,
195/// and flags. It uses a cache to store compiled regular expressions for efficiency.
196///
197/// # Errors
198/// Returns an error if the input arrays have mismatched lengths or if the regular expression fails to compile.
199fn 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            &regex_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            &regex_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            &regex_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            &regex_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            &regex_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            &regex_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    // A NULL in any argument produces a NULL result
296    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        // The defaults apply when the optional argument was not supplied.
317        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/// Compiles the patterns seen so far, keyed by `(pattern, flags)`.
328///
329/// Patterns are addressed by index rather than by reference so that `last` can
330/// memoize the previous row's pattern without holding a borrow of `indices`
331/// across rows. A literal pattern yields the same string on every row, so that
332/// memo means the common case never hashes a key.
333#[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
365/// Returns the 1-based character position of the `n`-th match of `pattern` in
366/// `value`, or 0 if there is no such match. The search begins at the 1-based
367/// character position `start`. A positive `subexpr` selects that capture group
368/// of the first match instead of the `n`-th match.
369fn 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    // A subexpression, when requested, takes precedence over the N-th match.
394    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        // `n` is 1-based, `nth` is 0-based.
401        pattern
402            .find_iter(search_slice)
403            .nth((n - 1) as usize)
404            .map(|matched| matched.start())
405    };
406
407    // Convert the byte offset within `search_slice` back to a 1-based character
408    // offset within `value`.
409    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        // let res_exp = re.unwrap();
486        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                // utf8
533                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                // let res_exp = re.unwrap();
538                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                // largeutf8
546                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                // utf8view
557                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                // utf8
579                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                // largeutf8
593                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                // utf8view
606                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                // utf8
631                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                // largeutf8
650                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                // utf8view
668                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            // utf8
707            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            // largeutf8
730            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            // utf8view
752            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        // A NULL in any argument produces a NULL result
777        let cases: Vec<Vec<ScalarValue>> = vec![
778            // NULL start
779            vec![
780                ScalarValue::Utf8(Some("abc".to_string())),
781                ScalarValue::Utf8(Some("b".to_string())),
782                ScalarValue::Int64(None),
783            ],
784            // NULL N
785            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            // NULL flags
792            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            // NULL subexpr
800            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}