diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index 94830ee360585..8fcc04862cb23 100644 --- a/datafusion/functions/Cargo.toml +++ b/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" diff --git a/datafusion/functions/benches/regexp_match.rs b/datafusion/functions/benches/regexp_match.rs new file mode 100644 index 0000000000000..d5929df07c81f --- /dev/null +++ b/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::() + }) + .collect::>() + .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::>(), + )) 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); diff --git a/datafusion/functions/src/regex/regexpmatch.rs b/datafusion/functions/src/regex/regexpmatch.rs index 34153d9c8ab96..918de5273b622 100644 --- a/datafusion/functions/src/regex/regexpmatch.rs +++ b/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 { 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::::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> { + 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 { 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:?}"); + } + } + } }