arrow_string/
regexp.rs

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

arrow_string/
regexp.rs
