datafusion_functions/regex/
regexpmatch.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//! Regex expressions
19use arrow::array::{Array, ArrayRef, AsArray, Datum};
20use arrow::compute::kernels::regexp;
21use arrow::datatypes::DataType;
22use arrow::datatypes::Field;
23use datafusion_common::Result;
24use datafusion_common::ScalarValue;
25use datafusion_common::exec_err;
26use datafusion_common::{arrow_datafusion_err, plan_err};
27use datafusion_expr::{ColumnarValue, Documentation, ScalarFunctionArgs, TypeSignature};
28use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
29use datafusion_macros::user_doc;
30use std::sync::Arc;
31
32#[user_doc(
33    doc_section(label = "Regular Expression Functions"),
34    description = "Returns the first [regular expression](https://docs.rs/regex/latest/regex/#syntax) matches in a string.",
35    syntax_example = "regexp_match(str, regexp[, flags])",
36    sql_example = r#"```sql
37            > select regexp_match('Köln', '[a-zA-Z]ö[a-zA-Z]{2}');
38            +---------------------------------------------------------+
39            | regexp_match(Utf8("Köln"),Utf8("[a-zA-Z]ö[a-zA-Z]{2}")) |
40            +---------------------------------------------------------+
41            | [Köln]                                                  |
42            +---------------------------------------------------------+
43            SELECT regexp_match('aBc', '(b|d)', 'i');
44            +---------------------------------------------------+
45            | regexp_match(Utf8("aBc"),Utf8("(b|d)"),Utf8("i")) |
46            +---------------------------------------------------+
47            | [B]                                               |
48            +---------------------------------------------------+
49```
50Additional examples can be found [here](https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/builtin_functions/regexp.rs)
51"#,
52    standard_argument(name = "str", prefix = "String"),
53    argument(
54        name = "regexp",
55        description = "Regular expression to match against.
56            Can be a constant, column, or function."
57    ),
58    argument(
59        name = "flags",
60        description = r#"Optional regular expression flags that control the behavior of the regular expression. Refer to the flags reference above for supported flags."#
61    )
62)]
63#[derive(Debug, PartialEq, Eq, Hash)]
64pub struct RegexpMatchFunc {
65    signature: Signature,
66}
67
68impl Default for RegexpMatchFunc {
69    fn default() -> Self {
70        Self::new()
71    }
72}
73
74impl RegexpMatchFunc {
75    pub fn new() -> Self {
76        use DataType::*;
77        Self {
78            signature: Signature::one_of(
79                vec![
80                    // Planner attempts coercion to the target type starting with the most preferred candidate.
81                    // For example, given input `(Utf8View, Utf8)`, it first tries coercing to `(Utf8View, Utf8View)`.
82                    // If that fails, it proceeds to `(Utf8, Utf8)`.
83                    TypeSignature::Exact(vec![Utf8View, Utf8View]),
84                    TypeSignature::Exact(vec![Utf8, Utf8]),
85                    TypeSignature::Exact(vec![LargeUtf8, LargeUtf8]),
86                    TypeSignature::Exact(vec![Utf8View, Utf8View, Utf8View]),
87                    TypeSignature::Exact(vec![Utf8, Utf8, Utf8]),
88                    TypeSignature::Exact(vec![LargeUtf8, LargeUtf8, LargeUtf8]),
89                ],
90                Volatility::Immutable,
91            ),
92        }
93    }
94}
95
96impl ScalarUDFImpl for RegexpMatchFunc {
97    fn name(&self) -> &str {
98        "regexp_match"
99    }
100
101    fn signature(&self) -> &Signature {
102        &self.signature
103    }
104
105    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
106        Ok(match &arg_types[0] {
107            DataType::Null => DataType::Null,
108            other => DataType::List(Arc::new(Field::new_list_field(other.clone(), true))),
109        })
110    }
111
112    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
113        let args = &args.args;
114
115        // A literal pattern is the common case, and handing it to the kernel as
116        // a scalar lets the regex be compiled once for the whole array. Any
117        // other argument shape falls through to the general path below.
118        if let Some(result) = regexp_match_scalar_pattern(args)? {
119            return Ok(ColumnarValue::Array(result));
120        }
121
122        let len = args
123            .iter()
124            .fold(Option::<usize>::None, |acc, arg| match arg {
125                ColumnarValue::Scalar(_) => acc,
126                ColumnarValue::Array(a) => Some(a.len()),
127            });
128
129        let is_scalar = len.is_none();
130        let inferred_length = len.unwrap_or(1);
131        let args = args
132            .iter()
133            .map(|arg| arg.to_array(inferred_length))
134            .collect::<Result<Vec<_>>>()?;
135
136        let result = regexp_match(&args);
137        if is_scalar {
138            // If all inputs are scalar, keeps output as scalar
139            let result = result.and_then(|arr| ScalarValue::try_from_array(&arr, 0));
140            result.map(ColumnarValue::Scalar)
141        } else {
142            result.map(ColumnarValue::Array)
143        }
144    }
145
146    fn documentation(&self) -> Option<&Documentation> {
147        self.doc()
148    }
149}
150
151/// Runs `regexp_match` with the pattern (and flags, if given) passed to the
152/// kernel as scalar [`Datum`]s, so the regex is compiled once for the whole
153/// array.
154///
155/// Applies when the values are an array, the pattern is a non-null scalar of
156/// the same string type as the values, and the flags, if given, are a scalar of
157/// that same type and are not the unsupported "global" flag.
158///
159/// Returns `Ok(None)` for every other argument shape, leaving the caller's
160/// general path to materialize each argument as an array, zip the rows, and
161/// raise whatever error the shape warrants.
162fn regexp_match_scalar_pattern(args: &[ColumnarValue]) -> Result<Option<ArrayRef>> {
163    let (values, pattern, flags) = match args {
164        [values, pattern] => (values, pattern, None),
165        [values, pattern, flags] => (values, pattern, Some(flags)),
166        _ => return Ok(None),
167    };
168
169    let (ColumnarValue::Array(values), ColumnarValue::Scalar(pattern)) =
170        (values, pattern)
171    else {
172        return Ok(None);
173    };
174    let flags = match flags {
175        // An array of flags has to be zipped with the values row by row.
176        Some(ColumnarValue::Array(_)) => return Ok(None),
177        Some(ColumnarValue::Scalar(flags)) => Some(flags),
178        None => None,
179    };
180
181    // The kernel requires the values, the pattern and the flags to share one
182    // string type.
183    let value_type = values.data_type();
184
185    if !matches!(pattern.try_as_str(), Some(Some(_)))
186        || &pattern.data_type() != value_type
187        || flags.is_some_and(|flags| {
188            flags.try_as_str() == Some(Some("g")) || &flags.data_type() != value_type
189        })
190    {
191        return Ok(None);
192    }
193
194    let pattern = pattern.to_scalar()?;
195    let flags = flags.map(ScalarValue::to_scalar).transpose()?;
196
197    regexp::regexp_match(
198        values,
199        &pattern,
200        flags.as_ref().map(|flags| flags as &dyn Datum),
201    )
202    .map(Some)
203    .map_err(|e| arrow_datafusion_err!(e))
204}
205
206pub fn regexp_match(args: &[ArrayRef]) -> Result<ArrayRef> {
207    match args.len() {
208        2 => regexp::regexp_match(&args[0], &args[1], None)
209            .map_err(|e| arrow_datafusion_err!(e)),
210        3 => {
211            match args[2].data_type() {
212                DataType::Utf8View => {
213                    if args[2].as_string_view().iter().any(|s| s == Some("g")) {
214                        return plan_err!(
215                            "regexp_match() does not support the \"global\" option"
216                        );
217                    }
218                }
219                DataType::Utf8 => {
220                    if args[2].as_string::<i32>().iter().any(|s| s == Some("g")) {
221                        return plan_err!(
222                            "regexp_match() does not support the \"global\" option"
223                        );
224                    }
225                }
226                DataType::LargeUtf8 => {
227                    if args[2].as_string::<i64>().iter().any(|s| s == Some("g")) {
228                        return plan_err!(
229                            "regexp_match() does not support the \"global\" option"
230                        );
231                    }
232                }
233                e => {
234                    return plan_err!(
235                        "regexp_match was called with unexpected data type {e:?}"
236                    );
237                }
238            }
239
240            regexp::regexp_match(&args[0], &args[1], Some(&args[2]))
241                .map_err(|e| arrow_datafusion_err!(e))
242        }
243        other => exec_err!(
244            "regexp_match was called with {other} arguments. It requires at least 2 and at most 3."
245        ),
246    }
247}
248#[cfg(test)]
249mod tests {
250    use crate::regex::regexpmatch::regexp_match;
251    use arrow::array::StringArray;
252    use arrow::array::{GenericStringBuilder, ListBuilder};
253    use std::sync::Arc;
254
255    #[test]
256    fn test_case_sensitive_regexp_match() {
257        let values = StringArray::from(vec!["abc"; 5]);
258        let patterns =
259            StringArray::from(vec!["^(a)", "^(A)", "(b|d)", "(B|D)", "^(b|c)"]);
260
261        let elem_builder: GenericStringBuilder<i32> = GenericStringBuilder::new();
262        let mut expected_builder = ListBuilder::new(elem_builder);
263        expected_builder.values().append_value("a");
264        expected_builder.append(true);
265        expected_builder.append(false);
266        expected_builder.values().append_value("b");
267        expected_builder.append(true);
268        expected_builder.append(false);
269        expected_builder.append(false);
270        let expected = expected_builder.finish();
271
272        let re = regexp_match(&[Arc::new(values), Arc::new(patterns)]).unwrap();
273
274        assert_eq!(re.as_ref(), &expected);
275    }
276
277    #[test]
278    fn test_case_insensitive_regexp_match() {
279        let values = StringArray::from(vec!["abc"; 5]);
280        let patterns =
281            StringArray::from(vec!["^(a)", "^(A)", "(b|d)", "(B|D)", "^(b|c)"]);
282        let flags = StringArray::from(vec!["i"; 5]);
283
284        let elem_builder: GenericStringBuilder<i32> = GenericStringBuilder::new();
285        let mut expected_builder = ListBuilder::new(elem_builder);
286        expected_builder.values().append_value("a");
287        expected_builder.append(true);
288        expected_builder.values().append_value("a");
289        expected_builder.append(true);
290        expected_builder.values().append_value("b");
291        expected_builder.append(true);
292        expected_builder.values().append_value("b");
293        expected_builder.append(true);
294        expected_builder.append(false);
295        let expected = expected_builder.finish();
296
297        let re = regexp_match(&[Arc::new(values), Arc::new(patterns), Arc::new(flags)])
298            .unwrap();
299
300        assert_eq!(re.as_ref(), &expected);
301    }
302
303    #[test]
304    fn test_unsupported_global_flag_regexp_match() {
305        let values = StringArray::from(vec!["abc"]);
306        let patterns = StringArray::from(vec!["^(a)"]);
307        let flags = StringArray::from(vec!["g"]);
308
309        let re_err =
310            regexp_match(&[Arc::new(values), Arc::new(patterns), Arc::new(flags)])
311                .expect_err("unsupported flag should have failed");
312
313        assert_eq!(
314            re_err.strip_backtrace(),
315            "Error during planning: regexp_match() does not support the \"global\" option"
316        );
317    }
318
319    /// The literal-pattern fast path must agree with the general path that
320    /// zips a pattern array with the values, for every argument shape.
321    #[test]
322    fn test_scalar_pattern_matches_array_pattern() {
323        use super::{RegexpMatchFunc, ScalarValue};
324        use arrow::array::{Array, ArrayRef};
325        use arrow::datatypes::{DataType, Field};
326        use datafusion_common::config::ConfigOptions;
327        use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
328
329        let values = Arc::new(StringArray::from(vec![
330            Some("abc"),
331            Some("ABC"),
332            None,
333            Some(""),
334            Some("a-b-c"),
335        ])) as ArrayRef;
336
337        for pattern in ["([a-z])(b)?", "^(A)", "no-match", "", "[a-z]+"] {
338            for flags in [None, Some("i")] {
339                let mut scalar_args = vec![
340                    ColumnarValue::Array(Arc::clone(&values)),
341                    ColumnarValue::Scalar(ScalarValue::Utf8(Some(pattern.to_string()))),
342                ];
343                let mut array_args = vec![
344                    Arc::clone(&values),
345                    Arc::new(StringArray::from(vec![pattern; values.len()])) as ArrayRef,
346                ];
347                if let Some(flags) = flags {
348                    scalar_args.push(ColumnarValue::Scalar(ScalarValue::Utf8(Some(
349                        flags.to_string(),
350                    ))));
351                    array_args
352                        .push(Arc::new(StringArray::from(vec![flags; values.len()]))
353                            as ArrayRef);
354                }
355
356                let arg_fields = scalar_args
357                    .iter()
358                    .enumerate()
359                    .map(|(idx, arg)| {
360                        Field::new(format!("arg_{idx}"), arg.data_type(), true).into()
361                    })
362                    .collect();
363                let actual = RegexpMatchFunc::new()
364                    .invoke_with_args(ScalarFunctionArgs {
365                        args: scalar_args,
366                        arg_fields,
367                        number_rows: values.len(),
368                        return_field: Field::new_list(
369                            "f",
370                            Field::new_list_field(DataType::Utf8, true),
371                            true,
372                        )
373                        .into(),
374                        config_options: Arc::new(ConfigOptions::default()),
375                    })
376                    .unwrap()
377                    .to_array(values.len())
378                    .unwrap();
379
380                let expected = regexp_match(&array_args).unwrap();
381                assert_eq!(&actual, &expected, "pattern={pattern:?} flags={flags:?}");
382            }
383        }
384    }
385}
datafusion_functions/regex/regexpmatch.rs

datafusion_functions/regex/
regexpmatch.rs
