perf: optimize regexp_match for literal pattern usage (20% faster)

datafusion/functions/Cargo.toml

@@ -153,6 +153,11 @@ harness = false
 name = "to_hex"
 required-features = ["string_expressions"]
 
+[[bench]]
+harness = false
+name = "regexp_match"
+required-features = ["regex_expressions"]
+
 [[bench]]
 harness = false
 name = "regx"

datafusion/functions/benches/regexp_match.rs

@@ -0,0 +1,137 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Benchmarks `regexp_match` through `invoke_with_args`, which is how a query
+//! plan calls it. The pattern (and flags) are literals, as in
+//! `regexp_match(col, '[a-z]+')`.
+
+use std::hint::black_box;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, StringArray};
+use arrow::compute::cast;
+use arrow::datatypes::{DataType, Field};
+use criterion::{Criterion, criterion_group, criterion_main};
+use datafusion_common::ScalarValue;
+use datafusion_common::config::ConfigOptions;
+use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
+use datafusion_functions::regex::regexpmatch::RegexpMatchFunc;
+use rand::Rng;
+use rand::distr::Alphanumeric;
+use rand::rngs::ThreadRng;
+
+const SIZE: usize = 1000;
+const PATTERN: &str = ".*([A-Z]{1}).*";
+
+fn data(rng: &mut ThreadRng) -> StringArray {
+    (0..SIZE)
+        .map(|_| {
+            rng.sample_iter(&Alphanumeric)
+                .take(7)
+                .map(char::from)
+                .collect::<String>()
+        })
+        .collect::<Vec<_>>()
+        .into()
+}
+
+fn run(c: &mut Criterion, name: &str, values: &ArrayRef, args: &[ColumnarValue]) {
+    let func = RegexpMatchFunc::new();
+    let arg_fields: Vec<_> = args
+        .iter()
+        .enumerate()
+        .map(|(idx, arg)| Field::new(format!("arg_{idx}"), arg.data_type(), true).into())
+        .collect();
+    let return_field = Arc::new(Field::new_list(
+        "f",
+        Field::new_list_field(values.data_type().clone(), true),
+        true,
+    ));
+    let config_options = Arc::new(ConfigOptions::default());
+
+    c.bench_function(name, |b| {
+        b.iter(|| {
+            black_box(
+                func.invoke_with_args(ScalarFunctionArgs {
+                    args: args.to_vec(),
+                    arg_fields: arg_fields.clone(),
+                    number_rows: SIZE,
+                    return_field: Arc::clone(&return_field),
+                    config_options: Arc::clone(&config_options),
+                })
+                .expect("regexp_match should work on valid values"),
+            )
+        })
+    });
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let mut rng = rand::rng();
+    let utf8 = Arc::new(data(&mut rng)) as ArrayRef;
+    let utf8view = cast(&utf8, &DataType::Utf8View).unwrap();
+
+    run(
+        c,
+        "regexp_match_1000 literal pattern",
+        &utf8,
+        &[
+            ColumnarValue::Array(Arc::clone(&utf8)),
+            ColumnarValue::Scalar(ScalarValue::Utf8(Some(PATTERN.to_string()))),
+        ],
+    );
+
+    run(
+        c,
+        "regexp_match_1000 literal pattern and flags",
+        &utf8,
+        &[
+            ColumnarValue::Array(Arc::clone(&utf8)),
+            ColumnarValue::Scalar(ScalarValue::Utf8(Some(PATTERN.to_string()))),
+            ColumnarValue::Scalar(ScalarValue::Utf8(Some("i".to_string()))),
+        ],
+    );
+
+    run(
+        c,
+        "regexp_match_1000 literal pattern utf8view",
+        &utf8view,
+        &[
+            ColumnarValue::Array(Arc::clone(&utf8view)),
+            ColumnarValue::Scalar(ScalarValue::Utf8View(Some(PATTERN.to_string()))),
+        ],
+    );
+
+    // Covers the path where the pattern varies per row and so cannot be
+    // compiled once for the whole array.
+    let patterns = Arc::new(StringArray::from(
+        (0..SIZE)
+            .map(|i| if i % 2 == 0 { PATTERN } else { "^(A).*" })
+            .collect::<Vec<_>>(),
+    )) as ArrayRef;
+    run(
+        c,
+        "regexp_match_1000 pattern array",
+        &utf8,
+        &[
+            ColumnarValue::Array(Arc::clone(&utf8)),
+            ColumnarValue::Array(patterns),
+        ],
+    );
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

datafusion/functions/src/regex/regexpmatch.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 //! Regex expressions
-use arrow::array::{Array, ArrayRef, AsArray};
+use arrow::array::{Array, ArrayRef, AsArray, Datum};
 use arrow::compute::kernels::regexp;
 use arrow::datatypes::DataType;
 use arrow::datatypes::Field;
@@ -116,6 +116,14 @@ impl ScalarUDFImpl for RegexpMatchFunc {
 
     fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
         let args = &args.args;
+
+        // A literal pattern is the common case, and handing it to the kernel as
+        // a scalar lets the regex be compiled once for the whole array. Any
+        // other argument shape falls through to the general path below.
+        if let Some(result) = regexp_match_scalar_pattern(args)? {
+            return Ok(ColumnarValue::Array(result));
+        }
+
         let len = args
             .iter()
             .fold(Option::<usize>::None, |acc, arg| match arg {
@@ -145,6 +153,61 @@ impl ScalarUDFImpl for RegexpMatchFunc {
     }
 }
 
+/// Runs `regexp_match` with the pattern (and flags, if given) passed to the
+/// kernel as scalar [`Datum`]s, so the regex is compiled once for the whole
+/// array.
+///
+/// Applies when the values are an array, the pattern is a non-null scalar of
+/// the same string type as the values, and the flags, if given, are a scalar of
+/// that same type and are not the unsupported "global" flag.
+///
+/// Returns `Ok(None)` for every other argument shape, leaving the caller's
+/// general path to materialize each argument as an array, zip the rows, and
+/// raise whatever error the shape warrants.
+fn regexp_match_scalar_pattern(args: &[ColumnarValue]) -> Result<Option<ArrayRef>> {
+    let (values, pattern, flags) = match args {
+        [values, pattern] => (values, pattern, None),
+        [values, pattern, flags] => (values, pattern, Some(flags)),
+        _ => return Ok(None),
+    };
+
+    let (ColumnarValue::Array(values), ColumnarValue::Scalar(pattern)) =
+        (values, pattern)
+    else {
+        return Ok(None);
+    };
+    let flags = match flags {
+        // An array of flags has to be zipped with the values row by row.
+        Some(ColumnarValue::Array(_)) => return Ok(None),
+        Some(ColumnarValue::Scalar(flags)) => Some(flags),
+        None => None,
+    };
+
+    // The kernel requires the values, the pattern and the flags to share one
+    // string type.
+    let value_type = values.data_type();
+
+    if !matches!(pattern.try_as_str(), Some(Some(_)))
+        || &pattern.data_type() != value_type
+        || flags.is_some_and(|flags| {
+            flags.try_as_str() == Some(Some("g")) || &flags.data_type() != value_type
+        })
+    {
+        return Ok(None);
+    }
+
+    let pattern = pattern.to_scalar()?;
+    let flags = flags.map(ScalarValue::to_scalar).transpose()?;
+
+    regexp::regexp_match(
+        values,
+        &pattern,
+        flags.as_ref().map(|flags| flags as &dyn Datum),
+    )
+    .map(Some)
+    .map_err(|e| arrow_datafusion_err!(e))
+}
+
 pub fn regexp_match(args: &[ArrayRef]) -> Result<ArrayRef> {
     match args.len() {
         2 => regexp::regexp_match(&args[0], &args[1], None)
@@ -257,4 +320,71 @@ mod tests {
             "Error during planning: regexp_match() does not support the \"global\" option"
         );
     }
+
+    /// The literal-pattern fast path must agree with the general path that
+    /// zips a pattern array with the values, for every argument shape.
+    #[test]
+    fn test_scalar_pattern_matches_array_pattern() {
+        use super::{RegexpMatchFunc, ScalarValue};
+        use arrow::array::{Array, ArrayRef};
+        use arrow::datatypes::{DataType, Field};
+        use datafusion_common::config::ConfigOptions;
+        use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
+
+        let values = Arc::new(StringArray::from(vec![
+            Some("abc"),
+            Some("ABC"),
+            None,
+            Some(""),
+            Some("a-b-c"),
+        ])) as ArrayRef;
+
+        for pattern in ["([a-z])(b)?", "^(A)", "no-match", "", "[a-z]+"] {
+            for flags in [None, Some("i")] {
+                let mut scalar_args = vec![
+                    ColumnarValue::Array(Arc::clone(&values)),
+                    ColumnarValue::Scalar(ScalarValue::Utf8(Some(pattern.to_string()))),
+                ];
+                let mut array_args = vec![
+                    Arc::clone(&values),
+                    Arc::new(StringArray::from(vec![pattern; values.len()])) as ArrayRef,
+                ];
+                if let Some(flags) = flags {
+                    scalar_args.push(ColumnarValue::Scalar(ScalarValue::Utf8(Some(
+                        flags.to_string(),
+                    ))));
+                    array_args
+                        .push(Arc::new(StringArray::from(vec![flags; values.len()]))
+                            as ArrayRef);
+                }
+
+                let arg_fields = scalar_args
+                    .iter()
+                    .enumerate()
+                    .map(|(idx, arg)| {
+                        Field::new(format!("arg_{idx}"), arg.data_type(), true).into()
+                    })
+                    .collect();
+                let actual = RegexpMatchFunc::new()
+                    .invoke_with_args(ScalarFunctionArgs {
+                        args: scalar_args,
+                        arg_fields,
+                        number_rows: values.len(),
+                        return_field: Field::new_list(
+                            "f",
+                            Field::new_list_field(DataType::Utf8, true),
+                            true,
+                        )
+                        .into(),
+                        config_options: Arc::new(ConfigOptions::default()),
+                    })
+                    .unwrap()
+                    .to_array(values.len())
+                    .unwrap();
+
+                let expected = regexp_match(&array_args).unwrap();
+                assert_eq!(&actual, &expected, "pattern={pattern:?} flags={flags:?}");
+            }
+        }
+    }
 }