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};
20use arrow::compute::kernels::regexp;
21use arrow::datatypes::DataType;
22use arrow::datatypes::Field;
23use datafusion_common::exec_err;
24use datafusion_common::ScalarValue;
25use datafusion_common::{arrow_datafusion_err, plan_err};
26use datafusion_common::{DataFusionError, Result};
27use datafusion_expr::{ColumnarValue, Documentation, TypeSignature};
28use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
29use datafusion_macros::user_doc;
30use std::any::Any;
31use std::sync::Arc;
32
33#[user_doc(
34    doc_section(label = "Regular Expression Functions"),
35    description = "Returns the first [regular expression](https://docs.rs/regex/latest/regex/#syntax) matches in a string.",
36    syntax_example = "regexp_match(str, regexp[, flags])",
37    sql_example = r#"```sql
38            > select regexp_match('Köln', '[a-zA-Z]ö[a-zA-Z]{2}');
39            +---------------------------------------------------------+
40            | regexp_match(Utf8("Köln"),Utf8("[a-zA-Z]ö[a-zA-Z]{2}")) |
41            +---------------------------------------------------------+
42            | [Köln]                                                  |
43            +---------------------------------------------------------+
44            SELECT regexp_match('aBc', '(b|d)', 'i');
45            +---------------------------------------------------+
46            | regexp_match(Utf8("aBc"),Utf8("(b|d)"),Utf8("i")) |
47            +---------------------------------------------------+
48            | [B]                                               |
49            +---------------------------------------------------+
50```
51Additional examples can be found [here](https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/regexp.rs)
52"#,
53    standard_argument(name = "str", prefix = "String"),
54    argument(
55        name = "regexp",
56        description = "Regular expression to match against.
57            Can be a constant, column, or function."
58    ),
59    argument(
60        name = "flags",
61        description = r#"Optional regular expression flags that control the behavior of the regular expression. The following flags are supported:
62  - **i**: case-insensitive: letters match both upper and lower case
63  - **m**: multi-line mode: ^ and $ match begin/end of line
64  - **s**: allow . to match \n
65  - **R**: enables CRLF mode: when multi-line mode is enabled, \r\n is used
66  - **U**: swap the meaning of x* and x*?"#
67    )
68)]
69#[derive(Debug)]
70pub struct RegexpMatchFunc {
71    signature: Signature,
72}
73
74impl Default for RegexpMatchFunc {
75    fn default() -> Self {
76        Self::new()
77    }
78}
79
80impl RegexpMatchFunc {
81    pub fn new() -> Self {
82        use DataType::*;
83        Self {
84            signature: Signature::one_of(
85                vec![
86                    // Planner attempts coercion to the target type starting with the most preferred candidate.
87                    // For example, given input `(Utf8View, Utf8)`, it first tries coercing to `(Utf8View, Utf8View)`.
88                    // If that fails, it proceeds to `(Utf8, Utf8)`.
89                    TypeSignature::Exact(vec![Utf8View, Utf8View]),
90                    TypeSignature::Exact(vec![Utf8, Utf8]),
91                    TypeSignature::Exact(vec![LargeUtf8, LargeUtf8]),
92                    TypeSignature::Exact(vec![Utf8View, Utf8View, Utf8View]),
93                    TypeSignature::Exact(vec![Utf8, Utf8, Utf8]),
94                    TypeSignature::Exact(vec![LargeUtf8, LargeUtf8, LargeUtf8]),
95                ],
96                Volatility::Immutable,
97            ),
98        }
99    }
100}
101
102impl ScalarUDFImpl for RegexpMatchFunc {
103    fn as_any(&self) -> &dyn Any {
104        self
105    }
106
107    fn name(&self) -> &str {
108        "regexp_match"
109    }
110
111    fn signature(&self) -> &Signature {
112        &self.signature
113    }
114
115    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
116        Ok(match &arg_types[0] {
117            DataType::Null => DataType::Null,
118            other => DataType::List(Arc::new(Field::new_list_field(other.clone(), true))),
119        })
120    }
121
122    fn invoke_with_args(
123        &self,
124        args: datafusion_expr::ScalarFunctionArgs,
125    ) -> Result<ColumnarValue> {
126        let args = &args.args;
127        let len = args
128            .iter()
129            .fold(Option::<usize>::None, |acc, arg| match arg {
130                ColumnarValue::Scalar(_) => acc,
131                ColumnarValue::Array(a) => Some(a.len()),
132            });
133
134        let is_scalar = len.is_none();
135        let inferred_length = len.unwrap_or(1);
136        let args = args
137            .iter()
138            .map(|arg| arg.to_array(inferred_length))
139            .collect::<Result<Vec<_>>>()?;
140
141        let result = regexp_match(&args);
142        if is_scalar {
143            // If all inputs are scalar, keeps output as scalar
144            let result = result.and_then(|arr| ScalarValue::try_from_array(&arr, 0));
145            result.map(ColumnarValue::Scalar)
146        } else {
147            result.map(ColumnarValue::Array)
148        }
149    }
150
151    fn documentation(&self) -> Option<&Documentation> {
152        self.doc()
153    }
154}
155
156pub fn regexp_match(args: &[ArrayRef]) -> Result<ArrayRef> {
157    match args.len() {
158        2 => {
159            regexp::regexp_match(&args[0], &args[1], None)
160                .map_err(|e| arrow_datafusion_err!(e))
161        }
162        3 => {
163            match args[2].data_type() {
164                DataType::Utf8View => {
165                    if args[2].as_string_view().iter().any(|s| s == Some("g")) {
166                        return plan_err!("regexp_match() does not support the \"global\" option");
167                    }
168                }
169                DataType::Utf8 => {
170                    if args[2].as_string::<i32>().iter().any(|s| s == Some("g")) {
171                        return plan_err!("regexp_match() does not support the \"global\" option");
172                    }
173                }
174                DataType::LargeUtf8 => {
175                    if args[2].as_string::<i64>().iter().any(|s| s == Some("g")) {
176                        return plan_err!("regexp_match() does not support the \"global\" option");
177                    }
178                }
179                e => {
180                    return plan_err!("regexp_match was called with unexpected data type {e:?}");
181                }
182            }
183
184            regexp::regexp_match(&args[0], &args[1], Some(&args[2]))
185                .map_err(|e| arrow_datafusion_err!(e))
186        }
187        other => exec_err!(
188            "regexp_match was called with {other} arguments. It requires at least 2 and at most 3."
189        ),
190    }
191}
192#[cfg(test)]
193mod tests {
194    use crate::regex::regexpmatch::regexp_match;
195    use arrow::array::StringArray;
196    use arrow::array::{GenericStringBuilder, ListBuilder};
197    use std::sync::Arc;
198
199    #[test]
200    fn test_case_sensitive_regexp_match() {
201        let values = StringArray::from(vec!["abc"; 5]);
202        let patterns =
203            StringArray::from(vec!["^(a)", "^(A)", "(b|d)", "(B|D)", "^(b|c)"]);
204
205        let elem_builder: GenericStringBuilder<i32> = GenericStringBuilder::new();
206        let mut expected_builder = ListBuilder::new(elem_builder);
207        expected_builder.values().append_value("a");
208        expected_builder.append(true);
209        expected_builder.append(false);
210        expected_builder.values().append_value("b");
211        expected_builder.append(true);
212        expected_builder.append(false);
213        expected_builder.append(false);
214        let expected = expected_builder.finish();
215
216        let re = regexp_match(&[Arc::new(values), Arc::new(patterns)]).unwrap();
217
218        assert_eq!(re.as_ref(), &expected);
219    }
220
221    #[test]
222    fn test_case_insensitive_regexp_match() {
223        let values = StringArray::from(vec!["abc"; 5]);
224        let patterns =
225            StringArray::from(vec!["^(a)", "^(A)", "(b|d)", "(B|D)", "^(b|c)"]);
226        let flags = StringArray::from(vec!["i"; 5]);
227
228        let elem_builder: GenericStringBuilder<i32> = GenericStringBuilder::new();
229        let mut expected_builder = ListBuilder::new(elem_builder);
230        expected_builder.values().append_value("a");
231        expected_builder.append(true);
232        expected_builder.values().append_value("a");
233        expected_builder.append(true);
234        expected_builder.values().append_value("b");
235        expected_builder.append(true);
236        expected_builder.values().append_value("b");
237        expected_builder.append(true);
238        expected_builder.append(false);
239        let expected = expected_builder.finish();
240
241        let re = regexp_match(&[Arc::new(values), Arc::new(patterns), Arc::new(flags)])
242            .unwrap();
243
244        assert_eq!(re.as_ref(), &expected);
245    }
246
247    #[test]
248    fn test_unsupported_global_flag_regexp_match() {
249        let values = StringArray::from(vec!["abc"]);
250        let patterns = StringArray::from(vec!["^(a)"]);
251        let flags = StringArray::from(vec!["g"]);
252
253        let re_err =
254            regexp_match(&[Arc::new(values), Arc::new(patterns), Arc::new(flags)])
255                .expect_err("unsupported flag should have failed");
256
257        assert_eq!(re_err.strip_backtrace(), "Error during planning: regexp_match() does not support the \"global\" option");
258    }
259}
datafusion_functions/regex/regexpmatch.rs

datafusion_functions/regex/
regexpmatch.rs
