This is an automated email from the ASF dual-hosted git repository.
dheres pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/master by this push:
new 15c19c314 Optimize `regex_replace` for scalar patterns (#3614)
15c19c314 is described below
commit 15c19c3140f4a03b169416ac88074065da5946e7
Author: Batuhan Taskaya <[email protected]>
AuthorDate: Tue Sep 27 09:32:42 2022 +0300
Optimize `regex_replace` for scalar patterns (#3614)
* Optimize `regex_replace` for scalar patterns
* Change the hot-path on `regexp_replace` to only variadic source (#2)
---
datafusion/physical-expr/src/functions.rs | 14 +-
datafusion/physical-expr/src/regex_expressions.rs | 288 ++++++++++++++++++++--
2 files changed, 282 insertions(+), 20 deletions(-)
diff --git a/datafusion/physical-expr/src/functions.rs
b/datafusion/physical-expr/src/functions.rs
index a34c1ab90..6997adc46 100644
--- a/datafusion/physical-expr/src/functions.rs
+++ b/datafusion/physical-expr/src/functions.rs
@@ -500,20 +500,22 @@ pub fn create_physical_fun(
BuiltinScalarFunction::RegexpReplace => {
Arc::new(|args| match args[0].data_type() {
DataType::Utf8 => {
- let func = invoke_if_regex_expressions_feature_flag!(
- regexp_replace,
+ let specializer_func =
invoke_if_regex_expressions_feature_flag!(
+ specialize_regexp_replace,
i32,
"regexp_replace"
);
- make_scalar_function(func)(args)
+ let func = specializer_func(args)?;
+ func(args)
}
DataType::LargeUtf8 => {
- let func = invoke_if_regex_expressions_feature_flag!(
- regexp_replace,
+ let specializer_func =
invoke_if_regex_expressions_feature_flag!(
+ specialize_regexp_replace,
i64,
"regexp_replace"
);
- make_scalar_function(func)(args)
+ let func = specializer_func(args)?;
+ func(args)
}
other => Err(DataFusionError::Internal(format!(
"Unsupported data type {:?} for function regexp_replace",
diff --git a/datafusion/physical-expr/src/regex_expressions.rs
b/datafusion/physical-expr/src/regex_expressions.rs
index c09c3c265..89805cade 100644
--- a/datafusion/physical-expr/src/regex_expressions.rs
+++ b/datafusion/physical-expr/src/regex_expressions.rs
@@ -21,16 +21,32 @@
//! Regex expressions
-use arrow::array::{ArrayRef, GenericStringArray, OffsetSizeTrait};
+use arrow::array::{
+ new_null_array, Array, ArrayRef, GenericStringArray, OffsetSizeTrait,
+};
use arrow::compute;
use datafusion_common::{DataFusionError, Result};
+use datafusion_expr::{ColumnarValue, ScalarFunctionImplementation};
use hashbrown::HashMap;
use lazy_static::lazy_static;
use regex::Regex;
use std::any::type_name;
use std::sync::Arc;
-macro_rules! downcast_string_arg {
+use crate::functions::make_scalar_function;
+
+macro_rules! fetch_string_arg {
+ ($ARG:expr, $NAME:expr, $T:ident, $EARLY_ABORT:ident) => {{
+ let array = downcast_string_array_arg!($ARG, $NAME, $T);
+ if array.is_null(0) {
+ return $EARLY_ABORT(array);
+ } else {
+ array.value(0)
+ }
+ }};
+}
+
+macro_rules! downcast_string_array_arg {
($ARG:expr, $NAME:expr, $T:ident) => {{
$ARG.as_any()
.downcast_ref::<GenericStringArray<T>>()
@@ -48,14 +64,14 @@ macro_rules! downcast_string_arg {
pub fn regexp_match<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef>
{
match args.len() {
2 => {
- let values = downcast_string_arg!(args[0], "string", T);
- let regex = downcast_string_arg!(args[1], "pattern", T);
+ let values = downcast_string_array_arg!(args[0], "string", T);
+ let regex = downcast_string_array_arg!(args[1], "pattern", T);
compute::regexp_match(values, regex,
None).map_err(DataFusionError::ArrowError)
}
3 => {
- let values = downcast_string_arg!(args[0], "string", T);
- let regex = downcast_string_arg!(args[1], "pattern", T);
- let flags = Some(downcast_string_arg!(args[2], "flags", T));
+ let values = downcast_string_array_arg!(args[0], "string", T);
+ let regex = downcast_string_array_arg!(args[1], "pattern", T);
+ let flags = Some(downcast_string_array_arg!(args[2], "flags", T));
compute::regexp_match(values, regex,
flags).map_err(DataFusionError::ArrowError)
}
other => Err(DataFusionError::Internal(format!(
@@ -80,14 +96,17 @@ fn regex_replace_posix_groups(replacement: &str) -> String {
///
/// example: `regexp_replace('Thomas', '.[mN]a.', 'M') = 'ThM'`
pub fn regexp_replace<T: OffsetSizeTrait>(args: &[ArrayRef]) ->
Result<ArrayRef> {
+ // Default implementation for regexp_replace, assumes all args are arrays
+ // and args is a sequence of 3 or 4 elements.
+
// creating Regex is expensive so create hashmap for memoization
let mut patterns: HashMap<String, Regex> = HashMap::new();
match args.len() {
3 => {
- let string_array = downcast_string_arg!(args[0], "string", T);
- let pattern_array = downcast_string_arg!(args[1], "pattern", T);
- let replacement_array = downcast_string_arg!(args[2],
"replacement", T);
+ let string_array = downcast_string_array_arg!(args[0], "string",
T);
+ let pattern_array = downcast_string_array_arg!(args[1], "pattern",
T);
+ let replacement_array = downcast_string_array_arg!(args[2],
"replacement", T);
let result = string_array
.iter()
@@ -120,10 +139,10 @@ pub fn regexp_replace<T: OffsetSizeTrait>(args:
&[ArrayRef]) -> Result<ArrayRef>
Ok(Arc::new(result) as ArrayRef)
}
4 => {
- let string_array = downcast_string_arg!(args[0], "string", T);
- let pattern_array = downcast_string_arg!(args[1], "pattern", T);
- let replacement_array = downcast_string_arg!(args[2],
"replacement", T);
- let flags_array = downcast_string_arg!(args[3], "flags", T);
+ let string_array = downcast_string_array_arg!(args[0], "string",
T);
+ let pattern_array = downcast_string_array_arg!(args[1], "pattern",
T);
+ let replacement_array = downcast_string_array_arg!(args[2],
"replacement", T);
+ let flags_array = downcast_string_array_arg!(args[3], "flags", T);
let result = string_array
.iter()
@@ -178,10 +197,125 @@ pub fn regexp_replace<T: OffsetSizeTrait>(args:
&[ArrayRef]) -> Result<ArrayRef>
}
}
+fn _regexp_replace_early_abort<T: OffsetSizeTrait>(
+ input_array: &GenericStringArray<T>,
+) -> Result<ArrayRef> {
+ // Mimicking the existing behavior of regexp_replace, if any of the scalar
arguments
+ // are actuall null, then the result will be an array of the same size but
with nulls.
+ Ok(new_null_array(input_array.data_type(), input_array.len()))
+}
+
+/// Special cased regex_replace implementation for the scenerio where
+/// the pattern, replacement and flags are static (arrays that are derived
+/// from scalars). This means we can skip regex caching system and basically
+/// hold a single Regex object for the replace operation. This also speeds
+/// up the pre-processing time of the replacement string, since it only
+/// needs to processed once.
+fn _regexp_replace_static_pattern_replace<T: OffsetSizeTrait>(
+ args: &[ArrayRef],
+) -> Result<ArrayRef> {
+ let string_array = downcast_string_array_arg!(args[0], "string", T);
+ let pattern = fetch_string_arg!(args[1], "pattern", T,
_regexp_replace_early_abort);
+ let replacement =
+ fetch_string_arg!(args[2], "replacement", T,
_regexp_replace_early_abort);
+ let flags = match args.len() {
+ 3 => None,
+ 4 => Some(fetch_string_arg!(args[3], "flags", T,
_regexp_replace_early_abort)),
+ other => {
+ return Err(DataFusionError::Internal(format!(
+ "regexp_replace was called with {} arguments. It requires at
least 3 and at most 4.",
+ other
+ )))
+ }
+ };
+
+ // Embed the flag (if it exists) into the pattern. Limit will determine
+ // whether this is a global match (as in replace all) or just a single
+ // replace operation.
+ let (pattern, limit) = match flags {
+ Some("g") => (pattern.to_string(), 0),
+ Some(flags) => (
+ format!("(?{}){}", flags.to_string().replace('g', ""), pattern),
+ !flags.contains('g') as usize,
+ ),
+ None => (pattern.to_string(), 1),
+ };
+
+ let re = Regex::new(&pattern)
+ .map_err(|err| DataFusionError::Execution(err.to_string()))?;
+
+ // Replaces the posix groups in the replacement string
+ // with rust ones.
+ let replacement = regex_replace_posix_groups(replacement);
+
+ let result = string_array
+ .iter()
+ .map(|string| {
+ string.map(|string| re.replacen(string, limit,
replacement.as_str()))
+ })
+ .collect::<GenericStringArray<T>>();
+ Ok(Arc::new(result) as ArrayRef)
+}
+
+/// Determine which implementation of the regexp_replace to use based
+/// on the given set of arguments.
+pub fn specialize_regexp_replace<T: OffsetSizeTrait>(
+ args: &[ColumnarValue],
+) -> Result<ScalarFunctionImplementation> {
+ // This will serve as a dispatch table where we can
+ // leverage it in order to determine whether the scalarity
+ // of the given set of arguments fits a better specialized
+ // function.
+ let (is_source_scalar, is_pattern_scalar, is_replacement_scalar,
is_flags_scalar) = (
+ matches!(args[0], ColumnarValue::Scalar(_)),
+ matches!(args[1], ColumnarValue::Scalar(_)),
+ matches!(args[2], ColumnarValue::Scalar(_)),
+ // The forth argument (flags) is optional; so in the event that
+ // it is not available, we'll claim that it is scalar.
+ matches!(args.get(3), Some(ColumnarValue::Scalar(_)) | None),
+ );
+
+ match (
+ is_source_scalar,
+ is_pattern_scalar,
+ is_replacement_scalar,
+ is_flags_scalar,
+ ) {
+ // This represents a very hot path for the case where the there is
+ // a single pattern that is being matched against and a single
replacement.
+ // This is extremely important to specialize on since it removes the
overhead
+ // of DF's in-house regex pattern cache (since there will be at most a
single
+ // pattern) and the pre-processing of the same replacement pattern at
each
+ // query.
+ //
+ // The flags needs to be a scalar as well since each pattern is
actually
+ // constructed with the flags embedded into the pattern itself. This
means
+ // even if the pattern itself is scalar, if the flags are an array then
+ // we will create many regexes and it is best to use the implementation
+ // that caches it. If there are no flags, we can simply ignore it here,
+ // and let the specialized function handle it.
+ (_, true, true, true) => {
+ // We still don't know the scalarity of source, so we need the
adapter
+ // even if it will do some extra work for the pattern and the
flags.
+ //
+ // TODO: maybe we need a way of telling the adapter on which
arguments
+ // it can skip filling (so that we won't create N - 1 redundant
cols).
+ Ok(make_scalar_function(
+ _regexp_replace_static_pattern_replace::<T>,
+ ))
+ }
+
+ // If there are no specialized implementations, we'll fall back to the
+ // generic implementation.
+ (_, _, _, _) => Ok(make_scalar_function(regexp_replace::<T>)),
+ }
+}
+
#[cfg(test)]
mod tests {
use super::*;
use arrow::array::*;
+ use datafusion_common::ScalarValue;
#[test]
fn test_case_sensitive_regexp_match() {
@@ -231,4 +365,130 @@ mod tests {
assert_eq!(re.as_ref(), &expected);
}
+
+ #[test]
+ fn test_static_pattern_regexp_replace() {
+ let values = StringArray::from(vec!["abc"; 5]);
+ let patterns = StringArray::from(vec!["b"; 5]);
+ let replacements = StringArray::from(vec!["foo"; 5]);
+ let expected = StringArray::from(vec!["afooc"; 5]);
+
+ let re = _regexp_replace_static_pattern_replace::<i32>(&[
+ Arc::new(values),
+ Arc::new(patterns),
+ Arc::new(replacements),
+ ])
+ .unwrap();
+
+ assert_eq!(re.as_ref(), &expected);
+ }
+
+ #[test]
+ fn test_static_pattern_regexp_replace_with_flags() {
+ let values = StringArray::from(vec!["abc", "ABC", "aBc", "AbC",
"aBC"]);
+ let patterns = StringArray::from(vec!["b"; 5]);
+ let replacements = StringArray::from(vec!["foo"; 5]);
+ let flags = StringArray::from(vec!["i"; 5]);
+ let expected =
+ StringArray::from(vec!["afooc", "AfooC", "afooc", "AfooC",
"afooC"]);
+
+ let re = _regexp_replace_static_pattern_replace::<i32>(&[
+ Arc::new(values),
+ Arc::new(patterns),
+ Arc::new(replacements),
+ Arc::new(flags),
+ ])
+ .unwrap();
+
+ assert_eq!(re.as_ref(), &expected);
+ }
+
+ #[test]
+ fn test_static_pattern_regexp_replace_early_abort() {
+ let values = StringArray::from(vec!["abc"; 5]);
+ let patterns = StringArray::from(vec![None; 5]);
+ let replacements = StringArray::from(vec!["foo"; 5]);
+ let expected = StringArray::from(vec![None; 5]);
+
+ let re = _regexp_replace_static_pattern_replace::<i32>(&[
+ Arc::new(values),
+ Arc::new(patterns),
+ Arc::new(replacements),
+ ])
+ .unwrap();
+
+ assert_eq!(re.as_ref(), &expected);
+ }
+
+ #[test]
+ fn test_static_pattern_regexp_replace_early_abort_flags() {
+ let values = StringArray::from(vec!["abc"; 5]);
+ let patterns = StringArray::from(vec!["a"; 5]);
+ let replacements = StringArray::from(vec!["foo"; 5]);
+ let flags = StringArray::from(vec![None; 5]);
+ let expected = StringArray::from(vec![None; 5]);
+
+ let re = _regexp_replace_static_pattern_replace::<i32>(&[
+ Arc::new(values),
+ Arc::new(patterns),
+ Arc::new(replacements),
+ Arc::new(flags),
+ ])
+ .unwrap();
+
+ assert_eq!(re.as_ref(), &expected);
+ }
+
+ #[test]
+ fn test_static_pattern_regexp_replace_pattern_error() {
+ let values = StringArray::from(vec!["abc"; 5]);
+ // Delibaretely using an invalid pattern to see how the single pattern
+ // error is propagated on regexp_replace.
+ let patterns = StringArray::from(vec!["["; 5]);
+ let replacements = StringArray::from(vec!["foo"; 5]);
+
+ let re = _regexp_replace_static_pattern_replace::<i32>(&[
+ Arc::new(values),
+ Arc::new(patterns),
+ Arc::new(replacements),
+ ]);
+ let pattern_err = re.expect_err("broken pattern should have failed");
+ assert_eq!(
+ pattern_err.to_string(),
+ "Execution error: regex parse error:\n [\n ^\nerror:
unclosed character class"
+ );
+ }
+
+ #[test]
+ fn test_regexp_can_specialize_all_cases() {
+ macro_rules! make_scalar {
+ () => {
+
ColumnarValue::Scalar(ScalarValue::Utf8(Some("foo".to_string())))
+ };
+ }
+
+ macro_rules! make_array {
+ () => {
+ ColumnarValue::Array(
+ Arc::new(StringArray::from(vec!["bar"; 2])) as ArrayRef
+ )
+ };
+ }
+
+ for source in [make_scalar!(), make_array!()] {
+ for pattern in [make_scalar!(), make_array!()] {
+ for replacement in [make_scalar!(), make_array!()] {
+ for flags in [Some(make_scalar!()), Some(make_array!()),
None] {
+ let mut args =
+ vec![source.clone(), pattern.clone(),
replacement.clone()];
+ if let Some(flags) = flags {
+ args.push(flags.clone());
+ }
+ let regex_func =
specialize_regexp_replace::<i32>(&args);
+ assert!(regex_func.is_ok());
+ }
+ }
+ }
+ }
+ }
}
