arrow_string/
regexp.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
18//! Defines kernel to extract substrings based on a regular
19//! expression of a \[Large\]StringArray
20
21use crate::like::StringArrayType;
22
23use arrow_array::builder::{
24    BooleanBufferBuilder, GenericStringBuilder, ListBuilder, StringViewBuilder,
25};
26use arrow_array::cast::AsArray;
27use arrow_array::*;
28use arrow_buffer::{BooleanBuffer, NullBuffer};
29use arrow_data::ArrayDataBuilder;
30use arrow_schema::{ArrowError, DataType, Field};
31use regex::Regex;
32
33use std::collections::HashMap;
34use std::sync::Arc;
35
36/// Return BooleanArray indicating which strings in an array match an array of
37/// regular expressions.
38///
39/// This is equivalent to the SQL `array ~ regex_array`, supporting
40/// [`StringArray`] / [`LargeStringArray`] / [`StringViewArray`].
41///
42/// If `regex_array` element has an empty value, the corresponding result value is always true.
43///
44/// `flags_array` are optional [`StringArray`] / [`LargeStringArray`] / [`StringViewArray`] flag,
45/// which allow special search modes, such as case-insensitive and multi-line mode.
46/// See the documentation [here](https://docs.rs/regex/1.5.4/regex/#grouping-and-flags)
47/// for more information.
48///
49/// # See Also
50/// * [`regexp_is_match_scalar`] for matching a single regular expression against an array of strings
51/// * [`regexp_match`] for extracting groups from a string array based on a regular expression
52///
53/// # Example
54/// ```
55/// # use arrow_array::{StringArray, BooleanArray};
56/// # use arrow_string::regexp::regexp_is_match;
57/// // First array is the array of strings to match
58/// let array = StringArray::from(vec!["Foo", "Bar", "FooBar", "Baz"]);
59/// // Second array is the array of regular expressions to match against
60/// let regex_array = StringArray::from(vec!["^Foo", "^Foo", "Bar$", "Baz"]);
61/// // Third array is the array of flags to use for each regular expression, if desired
62/// // (the type must be provided to satisfy type inference for the third parameter)
63/// let flags_array: Option<&StringArray> = None;
64/// // The result is a BooleanArray indicating when each string in `array`
65/// // matches the corresponding regular expression in `regex_array`
66/// let result = regexp_is_match(&array, &regex_array, flags_array).unwrap();
67/// assert_eq!(result, BooleanArray::from(vec![true, false, true, true]));
68/// ```
69pub fn regexp_is_match<'a, S1, S2, S3>(
70    array: &'a S1,
71    regex_array: &'a S2,
72    flags_array: Option<&'a S3>,
73) -> Result<BooleanArray, ArrowError>
74where
75    &'a S1: StringArrayType<'a>,
76    &'a S2: StringArrayType<'a>,
77    &'a S3: StringArrayType<'a>,
78{
79    if array.len() != regex_array.len() {
80        return Err(ArrowError::ComputeError(
81            "Cannot perform comparison operation on arrays of different length".to_string(),
82        ));
83    }
84
85    let nulls = NullBuffer::union(array.nulls(), regex_array.nulls());
86
87    let mut patterns: HashMap<String, Regex> = HashMap::new();
88    let mut result = BooleanBufferBuilder::new(array.len());
89
90    let complete_pattern = match flags_array {
91        Some(flags) => Box::new(
92            regex_array
93                .iter()
94                .zip(flags.iter())
95                .map(|(pattern, flags)| {
96                    pattern.map(|pattern| match flags {
97                        Some(flag) => format!("(?{flag}){pattern}"),
98                        None => pattern.to_string(),
99                    })
100                }),
101        ) as Box<dyn Iterator<Item = Option<String>>>,
102        None => Box::new(
103            regex_array
104                .iter()
105                .map(|pattern| pattern.map(|pattern| pattern.to_string())),
106        ),
107    };
108
109    array
110        .iter()
111        .zip(complete_pattern)
112        .map(|(value, pattern)| {
113            match (value, pattern) {
114                // Required for Postgres compatibility:
115                // SELECT 'foobarbequebaz' ~ ''); = true
116                (Some(_), Some(pattern)) if pattern == *"" => {
117                    result.append(true);
118                }
119                (Some(value), Some(pattern)) => {
120                    let existing_pattern = patterns.get(&pattern);
121                    let re = match existing_pattern {
122                        Some(re) => re,
123                        None => {
124                            let re = Regex::new(pattern.as_str()).map_err(|e| {
125                                ArrowError::ComputeError(format!(
126                                    "Regular expression did not compile: {e:?}"
127                                ))
128                            })?;
129                            patterns.entry(pattern).or_insert(re)
130                        }
131                    };
132                    result.append(re.is_match(value));
133                }
134                _ => result.append(false),
135            }
136            Ok(())
137        })
138        .collect::<Result<Vec<()>, ArrowError>>()?;
139
140    let data = unsafe {
141        ArrayDataBuilder::new(DataType::Boolean)
142            .len(array.len())
143            .buffers(vec![result.into()])
144            .nulls(nulls)
145            .build_unchecked()
146    };
147
148    Ok(BooleanArray::from(data))
149}
150
151/// Return BooleanArray indicating which strings in an array match a single regular expression.
152///
153/// This is equivalent to the SQL `array ~ regex_array`, supporting
154/// [`StringArray`] / [`LargeStringArray`] / [`StringViewArray`] and a scalar.
155///
156/// See the documentation on [`regexp_is_match`] for more details on arguments
157///
158/// # See Also
159/// * [`regexp_is_match`] for matching an array of regular expression against an array of strings
160/// * [`regexp_match`] for extracting groups from a string array based on a regular expression
161///
162/// # Example
163/// ```
164/// # use arrow_array::{StringArray, BooleanArray};
165/// # use arrow_string::regexp::regexp_is_match_scalar;
166/// // array of strings to match
167/// let array = StringArray::from(vec!["Foo", "Bar", "FooBar", "Baz"]);
168/// let regexp = "^Foo"; // regular expression to match against
169/// let flags: Option<&str> = None;  // flags can control the matching behavior
170/// // The result is a BooleanArray indicating when each string in `array`
171/// // matches the regular expression `regexp`
172/// let result = regexp_is_match_scalar(&array, regexp, None).unwrap();
173/// assert_eq!(result, BooleanArray::from(vec![true, false, true, false]));
174/// ```
175pub fn regexp_is_match_scalar<'a, S>(
176    array: &'a S,
177    regex: &str,
178    flag: Option<&str>,
179) -> Result<BooleanArray, ArrowError>
180where
181    &'a S: StringArrayType<'a>,
182{
183    let mut result = BooleanBufferBuilder::new(array.len());
184
185    let pattern = match flag {
186        Some(flag) => format!("(?{flag}){regex}"),
187        None => regex.to_string(),
188    };
189
190    if pattern.is_empty() {
191        result.append_n(array.len(), true);
192    } else {
193        let re = Regex::new(pattern.as_str()).map_err(|e| {
194            ArrowError::ComputeError(format!("Regular expression did not compile: {e:?}"))
195        })?;
196        for i in 0..array.len() {
197            let value = array.value(i);
198            result.append(re.is_match(value));
199        }
200    }
201
202    let values = BooleanBuffer::from(result);
203    let nulls = array
204        .nulls()
205        .map(|n| n.inner().sliced())
206        .and_then(|b| NullBuffer::from_unsliced_buffer(b, array.len()));
207    Ok(BooleanArray::new(values, nulls))
208}
209
210macro_rules! process_regexp_array_match {
211    ($array:expr, $regex_array:expr, $flags_array:expr, $list_builder:expr) => {
212        let mut patterns: HashMap<String, Regex> = HashMap::new();
213
214        let complete_pattern = match $flags_array {
215            Some(flags) => Box::new($regex_array.iter().zip(flags.iter()).map(
216                |(pattern, flags)| {
217                    pattern.map(|pattern| match flags {
218                        Some(value) => format!("(?{value}){pattern}"),
219                        None => pattern.to_string(),
220                    })
221                },
222            )) as Box<dyn Iterator<Item = Option<String>>>,
223            None => Box::new(
224                $regex_array
225                    .iter()
226                    .map(|pattern| pattern.map(|pattern| pattern.to_string())),
227            ),
228        };
229
230        $array
231            .iter()
232            .zip(complete_pattern)
233            .map(|(value, pattern)| {
234                match (value, pattern) {
235                    // Required for Postgres compatibility:
236                    // SELECT regexp_match('foobarbequebaz', ''); = {""}
237                    (Some(_), Some(pattern)) if pattern == *"" => {
238                        $list_builder.values().append_value("");
239                        $list_builder.append(true);
240                    }
241                    (Some(value), Some(pattern)) => {
242                        let existing_pattern = patterns.get(&pattern);
243                        let re = match existing_pattern {
244                            Some(re) => re,
245                            None => {
246                                let re = Regex::new(pattern.as_str()).map_err(|e| {
247                                    ArrowError::ComputeError(format!(
248                                        "Regular expression did not compile: {e:?}"
249                                    ))
250                                })?;
251                                patterns.entry(pattern).or_insert(re)
252                            }
253                        };
254                        match re.captures(value) {
255                            Some(caps) => {
256                                let mut iter = caps.iter();
257                                if caps.len() > 1 {
258                                    iter.next();
259                                }
260                                for m in iter.flatten() {
261                                    $list_builder.values().append_value(m.as_str());
262                                }
263
264                                $list_builder.append(true);
265                            }
266                            None => $list_builder.append(false),
267                        }
268                    }
269                    _ => $list_builder.append(false),
270                }
271                Ok(())
272            })
273            .collect::<Result<Vec<()>, ArrowError>>()?;
274    };
275}
276
277fn regexp_array_match<OffsetSize: OffsetSizeTrait>(
278    array: &GenericStringArray<OffsetSize>,
279    regex_array: &GenericStringArray<OffsetSize>,
280    flags_array: Option<&GenericStringArray<OffsetSize>>,
281) -> Result<ArrayRef, ArrowError> {
282    let builder: GenericStringBuilder<OffsetSize> = GenericStringBuilder::with_capacity(0, 0);
283    let mut list_builder = ListBuilder::new(builder);
284
285    process_regexp_array_match!(array, regex_array, flags_array, list_builder);
286
287    Ok(Arc::new(list_builder.finish()))
288}
289
290fn regexp_array_match_utf8view(
291    array: &StringViewArray,
292    regex_array: &StringViewArray,
293    flags_array: Option<&StringViewArray>,
294) -> Result<ArrayRef, ArrowError> {
295    let builder = StringViewBuilder::with_capacity(0);
296    let mut list_builder = ListBuilder::new(builder);
297
298    process_regexp_array_match!(array, regex_array, flags_array, list_builder);
299
300    Ok(Arc::new(list_builder.finish()))
301}
302
303fn get_scalar_pattern_flag<'a, OffsetSize: OffsetSizeTrait>(
304    regex_array: &'a dyn Array,
305    flag_array: Option<&'a dyn Array>,
306) -> (Option<&'a str>, Option<&'a str>) {
307    let regex = regex_array.as_string::<OffsetSize>();
308    let regex = regex.is_valid(0).then(|| regex.value(0));
309
310    if let Some(flag_array) = flag_array {
311        let flag = flag_array.as_string::<OffsetSize>();
312        (regex, flag.is_valid(0).then(|| flag.value(0)))
313    } else {
314        (regex, None)
315    }
316}
317
318fn get_scalar_pattern_flag_utf8view<'a>(
319    regex_array: &'a dyn Array,
320    flag_array: Option<&'a dyn Array>,
321) -> (Option<&'a str>, Option<&'a str>) {
322    let regex = regex_array.as_string_view();
323    let regex = regex.is_valid(0).then(|| regex.value(0));
324
325    if let Some(flag_array) = flag_array {
326        let flag = flag_array.as_string_view();
327        (regex, flag.is_valid(0).then(|| flag.value(0)))
328    } else {
329        (regex, None)
330    }
331}
332
333macro_rules! process_regexp_match {
334    ($array:expr, $regex:expr, $list_builder:expr) => {
335        // Reuse capture location storage across matches
336        let mut capture_locations = None;
337        $array
338            .iter()
339            .map(|value| {
340                match value {
341                    // Required for Postgres compatibility:
342                    // SELECT regexp_match('foobarbequebaz', ''); = {""}
343                    Some(_) if $regex.as_str().is_empty() => {
344                        $list_builder.values().append_value("");
345                        $list_builder.append(true);
346                    }
347                    Some(value) => {
348                        let locations =
349                            capture_locations.get_or_insert_with(|| $regex.capture_locations());
350                        match $regex.captures_read(locations, value) {
351                            Some(_) => {
352                                let mut groups = 0..locations.len();
353                                if locations.len() > 1 {
354                                    // Skip group 0 (the whole match) when explicit capture groups exist.
355                                    groups.next();
356                                }
357                                let iter = groups.map(|group| locations.get(group));
358                                for (start, end) in iter.flatten() {
359                                    $list_builder.values().append_value(&value[start..end]);
360                                }
361                                $list_builder.append(true);
362                            }
363                            None => $list_builder.append(false),
364                        }
365                    }
366                    None => $list_builder.append(false),
367                }
368                Ok(())
369            })
370            .collect::<Result<Vec<()>, ArrowError>>()?
371    };
372}
373
374fn regexp_scalar_match<OffsetSize: OffsetSizeTrait>(
375    array: &GenericStringArray<OffsetSize>,
376    regex: &Regex,
377) -> Result<ArrayRef, ArrowError> {
378    let builder: GenericStringBuilder<OffsetSize> = GenericStringBuilder::with_capacity(0, 0);
379    let mut list_builder = ListBuilder::new(builder);
380
381    process_regexp_match!(array, regex, list_builder);
382
383    Ok(Arc::new(list_builder.finish()))
384}
385
386fn regexp_scalar_match_utf8view(
387    array: &StringViewArray,
388    regex: &Regex,
389) -> Result<ArrayRef, ArrowError> {
390    let builder = StringViewBuilder::with_capacity(0);
391    let mut list_builder = ListBuilder::new(builder);
392
393    process_regexp_match!(array, regex, list_builder);
394
395    Ok(Arc::new(list_builder.finish()))
396}
397
398/// Extract all groups matched by a regular expression for a given String array.
399///
400/// Modelled after the Postgres [regexp_match].
401///
402/// Returns a ListArray of [`GenericStringArray`] with each element containing the leftmost-first
403/// match of the corresponding index in `regex_array` to string in `array`
404///
405/// If there is no match, the list element is NULL.
406///
407/// If a match is found, and the pattern contains no capturing parenthesized subexpressions,
408/// then the list element is a single-element [`GenericStringArray`] containing the substring
409/// matching the whole pattern.
410///
411/// If a match is found, and the pattern contains capturing parenthesized subexpressions, then the
412/// list element is a [`GenericStringArray`] whose n'th element is the substring matching
413/// the n'th capturing parenthesized subexpression of the pattern.
414///
415/// The flags parameter is an optional text string containing zero or more single-letter flags
416/// that change the function's behavior.
417///
418/// # See Also
419/// * [`regexp_is_match`] for matching (rather than extracting) a regular expression against an array of strings
420///
421/// [regexp_match]: https://www.postgresql.org/docs/current/functions-matching.html#FUNCTIONS-POSIX-REGEXP
422pub fn regexp_match(
423    array: &dyn Array,
424    regex_array: &dyn Datum,
425    flags_array: Option<&dyn Datum>,
426) -> Result<ArrayRef, ArrowError> {
427    let (rhs, is_rhs_scalar) = regex_array.get();
428
429    if array.data_type() != rhs.data_type() {
430        return Err(ArrowError::ComputeError(
431            "regexp_match() requires both array and pattern to be either Utf8, Utf8View or LargeUtf8"
432                .to_string(),
433        ));
434    }
435
436    let (flags, is_flags_scalar) = match flags_array {
437        Some(flags) => {
438            let (flags, is_flags_scalar) = flags.get();
439            (Some(flags), Some(is_flags_scalar))
440        }
441        None => (None, None),
442    };
443
444    if is_flags_scalar.is_some_and(|is_flags_scalar| is_rhs_scalar != is_flags_scalar) {
445        return Err(ArrowError::ComputeError(
446            "regexp_match() requires both pattern and flags to be either scalar or array"
447                .to_string(),
448        ));
449    }
450
451    if flags.is_some_and(|flags| rhs.data_type() != flags.data_type()) {
452        return Err(ArrowError::ComputeError(
453            "regexp_match() requires both pattern and flags to be either Utf8, Utf8View or LargeUtf8"
454                .to_string(),
455        ));
456    }
457
458    if is_rhs_scalar {
459        // Regex and flag is scalars
460        let (regex, flag) = match rhs.data_type() {
461            DataType::Utf8View => get_scalar_pattern_flag_utf8view(rhs, flags),
462            DataType::Utf8 => get_scalar_pattern_flag::<i32>(rhs, flags),
463            DataType::LargeUtf8 => get_scalar_pattern_flag::<i64>(rhs, flags),
464            _ => {
465                return Err(ArrowError::ComputeError(
466                    "regexp_match() requires pattern to be either Utf8, Utf8View or LargeUtf8"
467                        .to_string(),
468                ));
469            }
470        };
471
472        let Some(regex) = regex else {
473            return Ok(new_null_array(
474                &DataType::List(Arc::new(Field::new_list_field(
475                    array.data_type().clone(),
476                    true,
477                ))),
478                array.len(),
479            ));
480        };
481
482        let pattern = if let Some(flag) = flag {
483            format!("(?{flag}){regex}")
484        } else {
485            regex.to_string()
486        };
487
488        let re = Regex::new(pattern.as_str()).map_err(|e| {
489            ArrowError::ComputeError(format!("Regular expression did not compile: {e:?}"))
490        })?;
491
492        match array.data_type() {
493            DataType::Utf8View => regexp_scalar_match_utf8view(array.as_string_view(), &re),
494            DataType::Utf8 => regexp_scalar_match(array.as_string::<i32>(), &re),
495            DataType::LargeUtf8 => regexp_scalar_match(array.as_string::<i64>(), &re),
496            _ => Err(ArrowError::ComputeError(
497                "regexp_match() requires array to be either Utf8, Utf8View or LargeUtf8"
498                    .to_string(),
499            )),
500        }
501    } else {
502        match array.data_type() {
503            DataType::Utf8View => {
504                let regex_array = rhs.as_string_view();
505                let flags_array = flags.map(|flags| flags.as_string_view());
506                regexp_array_match_utf8view(array.as_string_view(), regex_array, flags_array)
507            }
508            DataType::Utf8 => {
509                let regex_array = rhs.as_string();
510                let flags_array = flags.map(|flags| flags.as_string());
511                regexp_array_match(array.as_string::<i32>(), regex_array, flags_array)
512            }
513            DataType::LargeUtf8 => {
514                let regex_array = rhs.as_string();
515                let flags_array = flags.map(|flags| flags.as_string());
516                regexp_array_match(array.as_string::<i64>(), regex_array, flags_array)
517            }
518            _ => Err(ArrowError::ComputeError(
519                "regexp_match() requires array to be either Utf8, Utf8View or LargeUtf8"
520                    .to_string(),
521            )),
522        }
523    }
524}
525
526#[cfg(test)]
527mod tests {
528    use super::*;
529
530    #[test]
531    fn match_scalar_pattern_capture_groups() {
532        macro_rules! check {
533            ($array_type:ty, $builder_type:ty) => {{
534                // "β" follows "aβ" to check that an optional group matched in
535                // one row is not reported for the next row.
536                let array = <$array_type>::from(vec![Some("aβ"), Some("β"), Some("none"), None]);
537                let cases: [(&str, [Option<&[&str]>; 4]); 4] = [
538                    // Optional groups that did not match are omitted.
539                    ("(a)?(β)", [Some(&["a", "β"]), Some(&["β"]), None, None]),
540                    // Without explicit groups, the whole match is returned.
541                    ("β", [Some(&["β"]), Some(&["β"]), None, None]),
542                    ("()", [Some(&[""]), Some(&[""]), Some(&[""]), None]),
543                    ("", [Some(&[""]), Some(&[""]), Some(&[""]), None]),
544                ];
545                for (pattern, rows) in cases {
546                    let mut expected = ListBuilder::new(<$builder_type>::new());
547                    for row in rows {
548                        match row {
549                            Some(values) => {
550                                for value in values {
551                                    expected.values().append_value(value);
552                                }
553                                expected.append(true);
554                            }
555                            None => expected.append(false),
556                        }
557                    }
558                    let expected = expected.finish();
559
560                    let pattern_scalar = Scalar::new(<$array_type>::from(vec![pattern]));
561                    let actual = regexp_match(&array, &pattern_scalar, None).unwrap();
562                    assert_eq!(actual.as_list::<i32>(), &expected, "{pattern}");
563                }
564            }};
565        }
566        check!(StringArray, GenericStringBuilder<i32>);
567        check!(LargeStringArray, GenericStringBuilder<i64>);
568        check!(StringViewArray, StringViewBuilder);
569    }
570
571    macro_rules! test_match_single_group {
572        ($test_name:ident, $values:expr, $patterns:expr, $arr_type:ty, $builder_type:ty, $expected:expr) => {
573            #[test]
574            fn $test_name() {
575                let array: $arr_type = <$arr_type>::from($values);
576                let pattern: $arr_type = <$arr_type>::from($patterns);
577
578                let actual = regexp_match(&array, &pattern, None).unwrap();
579
580                let elem_builder: $builder_type = <$builder_type>::new();
581                let mut expected_builder = ListBuilder::new(elem_builder);
582
583                for val in $expected {
584                    match val {
585                        Some(v) => {
586                            expected_builder.values().append_value(v);
587                            expected_builder.append(true);
588                        }
589                        None => expected_builder.append(false),
590                    }
591                }
592
593                let expected = expected_builder.finish();
594                let result = actual.as_any().downcast_ref::<ListArray>().unwrap();
595                assert_eq!(&expected, result);
596            }
597        };
598    }
599
600    test_match_single_group!(
601        match_single_group_string,
602        vec![
603            Some("abc-005-def"),
604            Some("X-7-5"),
605            Some("X545"),
606            None,
607            Some("foobarbequebaz"),
608            Some("foobarbequebaz"),
609        ],
610        vec![
611            r".*-(\d*)-.*",
612            r".*-(\d*)-.*",
613            r".*-(\d*)-.*",
614            r".*-(\d*)-.*",
615            r"(bar)(bequ1e)",
616            ""
617        ],
618        StringArray,
619        GenericStringBuilder<i32>,
620        [Some("005"), Some("7"), None, None, None, Some("")]
621    );
622    test_match_single_group!(
623        match_single_group_string_view,
624        vec![
625            Some("abc-005-def"),
626            Some("X-7-5"),
627            Some("X545"),
628            None,
629            Some("foobarbequebaz"),
630            Some("foobarbequebaz"),
631        ],
632        vec![
633            r".*-(\d*)-.*",
634            r".*-(\d*)-.*",
635            r".*-(\d*)-.*",
636            r".*-(\d*)-.*",
637            r"(bar)(bequ1e)",
638            ""
639        ],
640        StringViewArray,
641        StringViewBuilder,
642        [Some("005"), Some("7"), None, None, None, Some("")]
643    );
644
645    macro_rules! test_match_single_group_with_flags {
646        ($test_name:ident, $values:expr, $patterns:expr, $flags:expr, $array_type:ty, $builder_type:ty, $expected:expr) => {
647            #[test]
648            fn $test_name() {
649                let array: $array_type = <$array_type>::from($values);
650                let pattern: $array_type = <$array_type>::from($patterns);
651                let flags: $array_type = <$array_type>::from($flags);
652
653                let actual = regexp_match(&array, &pattern, Some(&flags)).unwrap();
654
655                let elem_builder: $builder_type = <$builder_type>::new();
656                let mut expected_builder = ListBuilder::new(elem_builder);
657
658                for val in $expected {
659                    match val {
660                        Some(v) => {
661                            expected_builder.values().append_value(v);
662                            expected_builder.append(true);
663                        }
664                        None => {
665                            expected_builder.append(false);
666                        }
667                    }
668                }
669
670                let expected = expected_builder.finish();
671                let result = actual.as_any().downcast_ref::<ListArray>().unwrap();
672                assert_eq!(&expected, result);
673            }
674        };
675    }
676
677    test_match_single_group_with_flags!(
678        match_single_group_with_flags_string,
679        vec![Some("abc-005-def"), Some("X-7-5"), Some("X545"), None],
680        vec![r"x.*-(\d*)-.*"; 4],
681        vec!["i"; 4],
682        StringArray,
683        GenericStringBuilder<i32>,
684        [None, Some("7"), None, None]
685    );
686    test_match_single_group_with_flags!(
687        match_single_group_with_flags_stringview,
688        vec![Some("abc-005-def"), Some("X-7-5"), Some("X545"), None],
689        vec![r"x.*-(\d*)-.*"; 4],
690        vec!["i"; 4],
691        StringViewArray,
692        StringViewBuilder,
693        [None, Some("7"), None, None]
694    );
695
696    macro_rules! test_match_scalar_pattern {
697        ($test_name:ident, $values:expr, $pattern:expr, $flag:expr, $array_type:ty, $builder_type:ty, $expected:expr) => {
698            #[test]
699            fn $test_name() {
700                let array: $array_type = <$array_type>::from($values);
701
702                let pattern_scalar = Scalar::new(<$array_type>::from(vec![$pattern; 1]));
703                let flag_scalar = Scalar::new(<$array_type>::from(vec![$flag; 1]));
704
705                let actual = regexp_match(&array, &pattern_scalar, Some(&flag_scalar)).unwrap();
706
707                let elem_builder: $builder_type = <$builder_type>::new();
708                let mut expected_builder = ListBuilder::new(elem_builder);
709
710                for val in $expected {
711                    match val {
712                        Some(v) => {
713                            expected_builder.values().append_value(v);
714                            expected_builder.append(true);
715                        }
716                        None => expected_builder.append(false),
717                    }
718                }
719
720                let expected = expected_builder.finish();
721                let result = actual.as_any().downcast_ref::<ListArray>().unwrap();
722                assert_eq!(&expected, result);
723            }
724        };
725    }
726
727    test_match_scalar_pattern!(
728        match_scalar_pattern_string_with_flags,
729        vec![
730            Some("abc-005-def"),
731            Some("x-7-5"),
732            Some("X-0-Y"),
733            Some("X545"),
734            None
735        ],
736        r"x.*-(\d*)-.*",
737        Some("i"),
738        StringArray,
739        GenericStringBuilder<i32>,
740        [None, Some("7"), Some("0"), None, None]
741    );
742    test_match_scalar_pattern!(
743        match_scalar_pattern_stringview_with_flags,
744        vec![
745            Some("abc-005-def"),
746            Some("x-7-5"),
747            Some("X-0-Y"),
748            Some("X545"),
749            None
750        ],
751        r"x.*-(\d*)-.*",
752        Some("i"),
753        StringViewArray,
754        StringViewBuilder,
755        [None, Some("7"), Some("0"), None, None]
756    );
757
758    test_match_scalar_pattern!(
759        match_scalar_pattern_string_no_flags,
760        vec![
761            Some("abc-005-def"),
762            Some("x-7-5"),
763            Some("X-0-Y"),
764            Some("X545"),
765            None
766        ],
767        r"x.*-(\d*)-.*",
768        None::<&str>,
769        StringArray,
770        GenericStringBuilder<i32>,
771        [None, Some("7"), None, None, None]
772    );
773    test_match_scalar_pattern!(
774        match_scalar_pattern_stringview_no_flags,
775        vec![
776            Some("abc-005-def"),
777            Some("x-7-5"),
778            Some("X-0-Y"),
779            Some("X545"),
780            None
781        ],
782        r"x.*-(\d*)-.*",
783        None::<&str>,
784        StringViewArray,
785        StringViewBuilder,
786        [None, Some("7"), None, None, None]
787    );
788
789    macro_rules! test_match_scalar_no_pattern {
790        ($test_name:ident, $values:expr, $array_type:ty, $pattern_type:expr, $builder_type:ty, $expected:expr) => {
791            #[test]
792            fn $test_name() {
793                let array: $array_type = <$array_type>::from($values);
794                let pattern = Scalar::new(new_null_array(&$pattern_type, 1));
795
796                let actual = regexp_match(&array, &pattern, None).unwrap();
797
798                let elem_builder: $builder_type = <$builder_type>::new();
799                let mut expected_builder = ListBuilder::new(elem_builder);
800
801                for val in $expected {
802                    match val {
803                        Some(v) => {
804                            expected_builder.values().append_value(v);
805                            expected_builder.append(true);
806                        }
807                        None => expected_builder.append(false),
808                    }
809                }
810
811                let expected = expected_builder.finish();
812                let result = actual.as_any().downcast_ref::<ListArray>().unwrap();
813                assert_eq!(&expected, result);
814            }
815        };
816    }
817
818    test_match_scalar_no_pattern!(
819        match_scalar_no_pattern_string,
820        vec![Some("abc-005-def"), Some("X-7-5"), Some("X545"), None],
821        StringArray,
822        DataType::Utf8,
823        GenericStringBuilder<i32>,
824        [None::<&str>, None, None, None]
825    );
826    test_match_scalar_no_pattern!(
827        match_scalar_no_pattern_stringview,
828        vec![Some("abc-005-def"), Some("X-7-5"), Some("X545"), None],
829        StringViewArray,
830        DataType::Utf8View,
831        StringViewBuilder,
832        [None::<&str>, None, None, None]
833    );
834
835    macro_rules! test_match_single_group_not_skip {
836        ($test_name:ident, $values:expr, $pattern:expr, $array_type:ty, $builder_type:ty, $expected:expr) => {
837            #[test]
838            fn $test_name() {
839                let array: $array_type = <$array_type>::from($values);
840                let pattern: $array_type = <$array_type>::from(vec![$pattern]);
841
842                let actual = regexp_match(&array, &pattern, None).unwrap();
843
844                let elem_builder: $builder_type = <$builder_type>::new();
845                let mut expected_builder = ListBuilder::new(elem_builder);
846
847                for val in $expected {
848                    match val {
849                        Some(v) => {
850                            expected_builder.values().append_value(v);
851                            expected_builder.append(true);
852                        }
853                        None => expected_builder.append(false),
854                    }
855                }
856
857                let expected = expected_builder.finish();
858                let result = actual.as_any().downcast_ref::<ListArray>().unwrap();
859                assert_eq!(&expected, result);
860            }
861        };
862    }
863
864    test_match_single_group_not_skip!(
865        match_single_group_not_skip_string,
866        vec![Some("foo"), Some("bar")],
867        r"foo",
868        StringArray,
869        GenericStringBuilder<i32>,
870        [Some("foo")]
871    );
872    test_match_single_group_not_skip!(
873        match_single_group_not_skip_stringview,
874        vec![Some("foo"), Some("bar")],
875        r"foo",
876        StringViewArray,
877        StringViewBuilder,
878        [Some("foo")]
879    );
880
881    macro_rules! test_flag_utf8 {
882        ($test_name:ident, $left:expr, $right:expr, $op:expr, $expected:expr) => {
883            #[test]
884            fn $test_name() {
885                let left = $left;
886                let right = $right;
887                let res = $op(&left, &right, None).unwrap();
888                let expected = $expected;
889                assert_eq!(expected.len(), res.len());
890                for i in 0..res.len() {
891                    let v = res.value(i);
892                    assert_eq!(v, expected[i]);
893                }
894            }
895        };
896        ($test_name:ident, $left:expr, $right:expr, $flag:expr, $op:expr, $expected:expr) => {
897            #[test]
898            fn $test_name() {
899                let left = $left;
900                let right = $right;
901                let flag = Some($flag);
902                let res = $op(&left, &right, flag.as_ref()).unwrap();
903                let expected = $expected;
904                assert_eq!(expected.len(), res.len());
905                for i in 0..res.len() {
906                    let v = res.value(i);
907                    assert_eq!(v, expected[i]);
908                }
909            }
910        };
911    }
912
913    macro_rules! test_flag_utf8_scalar {
914        ($test_name:ident, $left:expr, $right:expr, $op:expr, $expected:expr) => {
915            #[test]
916            fn $test_name() {
917                let left = $left;
918                let res = $op(&left, $right, None).unwrap();
919                let expected = $expected;
920                assert_eq!(expected.len(), res.len());
921                for i in 0..res.len() {
922                    let v = res.value(i);
923                    assert_eq!(
924                        v,
925                        expected[i],
926                        "unexpected result when comparing {} at position {} to {} ",
927                        left.value(i),
928                        i,
929                        $right
930                    );
931                }
932            }
933        };
934        ($test_name:ident, $left:expr, $right:expr, $flag:expr, $op:expr, $expected:expr) => {
935            #[test]
936            fn $test_name() {
937                let left = $left;
938                let flag = Some($flag);
939                let res = $op(&left, $right, flag).unwrap();
940                let expected = $expected;
941                assert_eq!(expected.len(), res.len());
942                for i in 0..res.len() {
943                    let v = res.value(i);
944                    assert_eq!(
945                        v,
946                        expected[i],
947                        "unexpected result when comparing {} at position {} to {} ",
948                        left.value(i),
949                        i,
950                        $right
951                    );
952                }
953            }
954        };
955    }
956
957    test_flag_utf8!(
958        test_array_regexp_is_match_utf8,
959        StringArray::from(vec!["arrow", "arrow", "arrow", "arrow", "arrow", "arrow"]),
960        StringArray::from(vec!["^ar", "^AR", "ow$", "OW$", "foo", ""]),
961        regexp_is_match::<StringArray, StringArray, StringArray>,
962        [true, false, true, false, false, true]
963    );
964    test_flag_utf8!(
965        test_array_regexp_is_match_utf8_insensitive,
966        StringArray::from(vec!["arrow", "arrow", "arrow", "arrow", "arrow", "arrow"]),
967        StringArray::from(vec!["^ar", "^AR", "ow$", "OW$", "foo", ""]),
968        StringArray::from(vec!["i"; 6]),
969        regexp_is_match,
970        [true, true, true, true, false, true]
971    );
972
973    test_flag_utf8_scalar!(
974        test_array_regexp_is_match_utf8_scalar,
975        StringArray::from(vec!["arrow", "ARROW", "parquet", "PARQUET"]),
976        "^ar",
977        regexp_is_match_scalar,
978        [true, false, false, false]
979    );
980    test_flag_utf8_scalar!(
981        test_array_regexp_is_match_utf8_scalar_empty,
982        StringArray::from(vec!["arrow", "ARROW", "parquet", "PARQUET"]),
983        "",
984        regexp_is_match_scalar,
985        [true, true, true, true]
986    );
987    test_flag_utf8_scalar!(
988        test_array_regexp_is_match_utf8_scalar_insensitive,
989        StringArray::from(vec!["arrow", "ARROW", "parquet", "PARQUET"]),
990        "^ar",
991        "i",
992        regexp_is_match_scalar,
993        [true, true, false, false]
994    );
995
996    test_flag_utf8!(
997        tes_array_regexp_is_match,
998        StringViewArray::from(vec!["arrow", "arrow", "arrow", "arrow", "arrow", "arrow"]),
999        StringViewArray::from(vec!["^ar", "^AR", "ow$", "OW$", "foo", ""]),
1000        regexp_is_match::<StringViewArray, StringViewArray, StringViewArray>,
1001        [true, false, true, false, false, true]
1002    );
1003    test_flag_utf8!(
1004        test_array_regexp_is_match_2,
1005        StringViewArray::from(vec!["arrow", "arrow", "arrow", "arrow", "arrow", "arrow"]),
1006        StringArray::from(vec!["^ar", "^AR", "ow$", "OW$", "foo", ""]),
1007        regexp_is_match::<StringViewArray, GenericStringArray<i32>, GenericStringArray<i32>>,
1008        [true, false, true, false, false, true]
1009    );
1010    test_flag_utf8!(
1011        test_array_regexp_is_match_insensitive,
1012        StringViewArray::from(vec![
1013            "Official Rust implementation of Apache Arrow",
1014            "apache/arrow-rs",
1015            "apache/arrow-rs",
1016            "parquet",
1017            "parquet",
1018            "row",
1019            "row",
1020        ]),
1021        StringViewArray::from(vec![
1022            ".*rust implement.*",
1023            "^ap",
1024            "^AP",
1025            "et$",
1026            "ET$",
1027            "foo",
1028            ""
1029        ]),
1030        StringViewArray::from(vec!["i"; 7]),
1031        regexp_is_match::<StringViewArray, StringViewArray, StringViewArray>,
1032        [true, true, true, true, true, false, true]
1033    );
1034    test_flag_utf8!(
1035        test_array_regexp_is_match_insensitive_2,
1036        LargeStringArray::from(vec!["arrow", "arrow", "arrow", "arrow", "arrow", "arrow"]),
1037        StringViewArray::from(vec!["^ar", "^AR", "ow$", "OW$", "foo", ""]),
1038        StringArray::from(vec!["i"; 6]),
1039        regexp_is_match::<GenericStringArray<i64>, StringViewArray, GenericStringArray<i32>>,
1040        [true, true, true, true, false, true]
1041    );
1042
1043    test_flag_utf8_scalar!(
1044        test_array_regexp_is_match_scalar,
1045        StringViewArray::from(vec![
1046            "apache/arrow-rs",
1047            "APACHE/ARROW-RS",
1048            "parquet",
1049            "PARQUET",
1050        ]),
1051        "^ap",
1052        regexp_is_match_scalar::<StringViewArray>,
1053        [true, false, false, false]
1054    );
1055    test_flag_utf8_scalar!(
1056        test_array_regexp_is_match_scalar_empty,
1057        StringViewArray::from(vec![
1058            "apache/arrow-rs",
1059            "APACHE/ARROW-RS",
1060            "parquet",
1061            "PARQUET",
1062        ]),
1063        "",
1064        regexp_is_match_scalar::<StringViewArray>,
1065        [true, true, true, true]
1066    );
1067    test_flag_utf8_scalar!(
1068        test_array_regexp_is_match_scalar_insensitive,
1069        StringViewArray::from(vec![
1070            "apache/arrow-rs",
1071            "APACHE/ARROW-RS",
1072            "parquet",
1073            "PARQUET",
1074        ]),
1075        "^ap",
1076        "i",
1077        regexp_is_match_scalar::<StringViewArray>,
1078        [true, true, false, false]
1079    );
1080}
arrow_string/regexp.rs

arrow_string/
regexp.rs
