This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 04e8e53366 Fix: Internal error in regexp_replace() for some StringView
input (#12203)
04e8e53366 is described below
commit 04e8e5336670f785ac239245bcfad280df571cd3
Author: WeblWabl <[email protected]>
AuthorDate: Thu Sep 12 11:05:55 2024 -0500
Fix: Internal error in regexp_replace() for some StringView input (#12203)
* Fix: Internal error in regexp_replace() for some StringView input
* fix regex bench
* fmt
* fix bench regx
* clippy
* fmt
* adds tests for flags + includes type signature for utf8view with flag
* fix: adding collect for string view type
---
datafusion/functions/benches/regx.rs | 14 +-
datafusion/functions/src/regex/regexpreplace.rs | 258 ++++++++++++++-------
datafusion/sqllogictest/test_files/string_view.slt | 90 +++++++
3 files changed, 268 insertions(+), 94 deletions(-)
diff --git a/datafusion/functions/benches/regx.rs
b/datafusion/functions/benches/regx.rs
index 23d57f38ef..45bfa23511 100644
--- a/datafusion/functions/benches/regx.rs
+++ b/datafusion/functions/benches/regx.rs
@@ -18,7 +18,7 @@
extern crate criterion;
use arrow::array::builder::StringBuilder;
-use arrow::array::{ArrayRef, StringArray};
+use arrow::array::{ArrayRef, AsArray, StringArray};
use criterion::{black_box, criterion_group, criterion_main, Criterion};
use datafusion_functions::regex::regexplike::regexp_like;
use datafusion_functions::regex::regexpmatch::regexp_match;
@@ -122,12 +122,12 @@ fn criterion_benchmark(c: &mut Criterion) {
b.iter(|| {
black_box(
- regexp_replace::<i32>(&[
- Arc::clone(&data),
- Arc::clone(®ex),
- Arc::clone(&replacement),
- Arc::clone(&flags),
- ])
+ regexp_replace::<i32, _, _>(
+ data.as_string::<i32>(),
+ regex.as_string::<i32>(),
+ replacement.as_string::<i32>(),
+ Some(&flags),
+ )
.expect("regexp_replace should work on valid values"),
)
})
diff --git a/datafusion/functions/src/regex/regexpreplace.rs
b/datafusion/functions/src/regex/regexpreplace.rs
index 0b0f7287e1..3eb72a1fb5 100644
--- a/datafusion/functions/src/regex/regexpreplace.rs
+++ b/datafusion/functions/src/regex/regexpreplace.rs
@@ -16,13 +16,13 @@
// under the License.
//! Regx expressions
-use arrow::array::new_null_array;
-use arrow::array::ArrayAccessor;
use arrow::array::ArrayDataBuilder;
use arrow::array::BufferBuilder;
use arrow::array::GenericStringArray;
use arrow::array::StringViewBuilder;
+use arrow::array::{new_null_array, ArrayIter, AsArray};
use arrow::array::{Array, ArrayRef, OffsetSizeTrait};
+use arrow::array::{ArrayAccessor, StringViewArray};
use arrow::datatypes::DataType;
use datafusion_common::cast::as_string_view_array;
use datafusion_common::exec_err;
@@ -59,6 +59,7 @@ impl RegexpReplaceFunc {
Exact(vec![Utf8, Utf8, Utf8]),
Exact(vec![Utf8View, Utf8, Utf8]),
Exact(vec![Utf8, Utf8, Utf8, Utf8]),
+ Exact(vec![Utf8View, Utf8, Utf8, Utf8]),
],
Volatility::Immutable,
),
@@ -187,104 +188,147 @@ fn regex_replace_posix_groups(replacement: &str) ->
String {
/// # Ok(())
/// # }
/// ```
-pub fn regexp_replace<T: OffsetSizeTrait>(args: &[ArrayRef]) ->
Result<ArrayRef> {
+pub fn regexp_replace<'a, T: OffsetSizeTrait, V, B>(
+ string_array: V,
+ pattern_array: B,
+ replacement_array: B,
+ flags: Option<&ArrayRef>,
+) -> Result<ArrayRef>
+where
+ V: ArrayAccessor<Item = &'a str>,
+ B: ArrayAccessor<Item = &'a str>,
+{
// 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 = as_generic_string_array::<T>(&args[0])?;
- let pattern_array = as_generic_string_array::<T>(&args[1])?;
- let replacement_array = as_generic_string_array::<T>(&args[2])?;
-
- let result = string_array
- .iter()
- .zip(pattern_array.iter())
- .zip(replacement_array.iter())
- .map(|((string, pattern), replacement)| match (string, pattern,
replacement) {
- (Some(string), Some(pattern), Some(replacement)) => {
- let replacement = regex_replace_posix_groups(replacement);
-
- // if patterns hashmap already has regexp then use else
create and return
- let re = match patterns.get(pattern) {
- Some(re) => Ok(re),
- None => {
- match Regex::new(pattern) {
- Ok(re) => {
- patterns.insert(pattern.to_string(), re);
- Ok(patterns.get(pattern).unwrap())
+ let datatype = string_array.data_type().to_owned();
+
+ let string_array_iter = ArrayIter::new(string_array);
+ let pattern_array_iter = ArrayIter::new(pattern_array);
+ let replacement_array_iter = ArrayIter::new(replacement_array);
+
+ match flags {
+ None => {
+ let result_iter = string_array_iter
+ .zip(pattern_array_iter)
+ .zip(replacement_array_iter)
+ .map(|((string, pattern), replacement)| {
+ match (string, pattern, replacement) {
+ (Some(string), Some(pattern), Some(replacement)) => {
+ let replacement =
regex_replace_posix_groups(replacement);
+ // if patterns hashmap already has regexp then use
else create and return
+ let re = match patterns.get(pattern) {
+ Some(re) => Ok(re),
+ None => match Regex::new(pattern) {
+ Ok(re) => {
+ patterns.insert(pattern.to_string(),
re);
+ Ok(patterns.get(pattern).unwrap())
+ }
+ Err(err) => {
+
Err(DataFusionError::External(Box::new(err)))
+ }
},
- Err(err) =>
Err(DataFusionError::External(Box::new(err))),
- }
- }
- };
+ };
- Some(re.map(|re| re.replace(string,
replacement.as_str()))).transpose()
+ Some(re.map(|re| re.replace(string,
replacement.as_str())))
+ .transpose()
+ }
+ _ => Ok(None),
+ }
+ });
+
+ match datatype {
+ DataType::Utf8 | DataType::LargeUtf8 => {
+ let result =
+
result_iter.collect::<Result<GenericStringArray<T>>>()?;
+ Ok(Arc::new(result) as ArrayRef)
}
- _ => Ok(None)
- })
- .collect::<Result<GenericStringArray<T>>>()?;
-
- Ok(Arc::new(result) as ArrayRef)
+ DataType::Utf8View => {
+ let result =
result_iter.collect::<Result<StringViewArray>>()?;
+ Ok(Arc::new(result) as ArrayRef)
+ }
+ other => {
+ exec_err!(
+ "Unsupported data type {other:?} for function
regex_replace"
+ )
+ }
+ }
}
- 4 => {
- let string_array = as_generic_string_array::<T>(&args[0])?;
- let pattern_array = as_generic_string_array::<T>(&args[1])?;
- let replacement_array = as_generic_string_array::<T>(&args[2])?;
- let flags_array = as_generic_string_array::<T>(&args[3])?;
-
- let result = string_array
- .iter()
- .zip(pattern_array.iter())
- .zip(replacement_array.iter())
- .zip(flags_array.iter())
- .map(|(((string, pattern), replacement), flags)| match (string,
pattern, replacement, flags) {
- (Some(string), Some(pattern), Some(replacement), Some(flags))
=> {
- let replacement = regex_replace_posix_groups(replacement);
-
- // format flags into rust pattern
- let (pattern, replace_all) = if flags == "g" {
- (pattern.to_string(), true)
- } else if flags.contains('g') {
- (format!("(?{}){}", flags.to_string().replace('g',
""), pattern), true)
- } else {
- (format!("(?{flags}){pattern}"), false)
- };
-
- // if patterns hashmap already has regexp then use else
create and return
- let re = match patterns.get(&pattern) {
- Some(re) => Ok(re),
- None => {
- match Regex::new(pattern.as_str()) {
- Ok(re) => {
- patterns.insert(pattern.clone(), re);
- Ok(patterns.get(&pattern).unwrap())
+ Some(flags) => {
+ let flags_array = as_generic_string_array::<T>(flags)?;
+
+ let result_iter = string_array_iter
+ .zip(pattern_array_iter)
+ .zip(replacement_array_iter)
+ .zip(flags_array.iter())
+ .map(|(((string, pattern), replacement), flags)| {
+ match (string, pattern, replacement, flags) {
+ (Some(string), Some(pattern), Some(replacement),
Some(flags)) => {
+ let replacement =
regex_replace_posix_groups(replacement);
+
+ // format flags into rust pattern
+ let (pattern, replace_all) = if flags == "g" {
+ (pattern.to_string(), true)
+ } else if flags.contains('g') {
+ (
+ format!(
+ "(?{}){}",
+ flags.to_string().replace('g', ""),
+ pattern
+ ),
+ true,
+ )
+ } else {
+ (format!("(?{flags}){pattern}"), false)
+ };
+
+ // if patterns hashmap already has regexp then use
else create and return
+ let re = match patterns.get(&pattern) {
+ Some(re) => Ok(re),
+ None => match Regex::new(pattern.as_str()) {
+ Ok(re) => {
+ patterns.insert(pattern.clone(), re);
+ Ok(patterns.get(&pattern).unwrap())
+ }
+ Err(err) => {
+
Err(DataFusionError::External(Box::new(err)))
+ }
},
- Err(err) =>
Err(DataFusionError::External(Box::new(err))),
- }
+ };
+
+ Some(re.map(|re| {
+ if replace_all {
+ re.replace_all(string,
replacement.as_str())
+ } else {
+ re.replace(string, replacement.as_str())
+ }
+ }))
+ .transpose()
}
- };
-
- Some(re.map(|re| {
- if replace_all {
- re.replace_all(string, replacement.as_str())
- } else {
- re.replace(string, replacement.as_str())
- }
- })).transpose()
+ _ => Ok(None),
+ }
+ });
+
+ match datatype {
+ DataType::Utf8 | DataType::LargeUtf8 => {
+ let result =
+
result_iter.collect::<Result<GenericStringArray<T>>>()?;
+ Ok(Arc::new(result) as ArrayRef)
}
- _ => Ok(None)
- })
- .collect::<Result<GenericStringArray<T>>>()?;
-
- Ok(Arc::new(result) as ArrayRef)
+ DataType::Utf8View => {
+ let result =
result_iter.collect::<Result<StringViewArray>>()?;
+ Ok(Arc::new(result) as ArrayRef)
+ }
+ other => {
+ exec_err!(
+ "Unsupported data type {other:?} for function
regex_replace"
+ )
+ }
+ }
}
- other => exec_err!(
- "regexp_replace was called with {other} arguments. It requires at
least 3 and at most 4."
- ),
}
}
@@ -495,7 +539,47 @@ pub fn specialize_regexp_replace<T: OffsetSizeTrait>(
.iter()
.map(|arg| arg.clone().into_array(inferred_length))
.collect::<Result<Vec<_>>>()?;
- regexp_replace::<T>(&args)
+
+ match args[0].data_type() {
+ DataType::Utf8View => {
+ let string_array = args[0].as_string_view();
+ let pattern_array = args[1].as_string::<i32>();
+ let replacement_array = args[2].as_string::<i32>();
+ regexp_replace::<i32, _, _>(
+ string_array,
+ pattern_array,
+ replacement_array,
+ args.get(3),
+ )
+ }
+ DataType::Utf8 => {
+ let string_array = args[0].as_string::<i32>();
+ let pattern_array = args[1].as_string::<i32>();
+ let replacement_array = args[2].as_string::<i32>();
+ regexp_replace::<i32, _, _>(
+ string_array,
+ pattern_array,
+ replacement_array,
+ args.get(3),
+ )
+ }
+ DataType::LargeUtf8 => {
+ let string_array = args[0].as_string::<i64>();
+ let pattern_array = args[1].as_string::<i64>();
+ let replacement_array = args[2].as_string::<i64>();
+ regexp_replace::<i64, _, _>(
+ string_array,
+ pattern_array,
+ replacement_array,
+ args.get(3),
+ )
+ }
+ other => {
+ exec_err!(
+ "Unsupported data type {other:?} for function
regex_replace"
+ )
+ }
+ }
}
}
}
diff --git a/datafusion/sqllogictest/test_files/string_view.slt
b/datafusion/sqllogictest/test_files/string_view.slt
index 171b8ec6c1..161e4e88de 100644
--- a/datafusion/sqllogictest/test_files/string_view.slt
+++ b/datafusion/sqllogictest/test_files/string_view.slt
@@ -459,6 +459,96 @@ Xiangpeng
Raphael
NULL
+### Test REGEXP_REPLACE
+
+# Should run REGEXP_REPLACE with Scalar value for utf8view
+query T
+SELECT
+ REGEXP_REPLACE(column1_utf8view, 'e', 'f') AS k
+FROM test;
+----
+Andrfw
+Xiangpfng
+Raphafl
+NULL
+
+# Should run REGEXP_REPLACE with Scalar value for utf8view with flag
+query T
+SELECT
+ REGEXP_REPLACE(column1_utf8view, 'e', 'f', 'i') AS k
+FROM test;
+----
+Andrfw
+Xiangpfng
+Raphafl
+NULL
+
+# Should run REGEXP_REPLACE with Scalar value for utf8
+query T
+SELECT
+ REGEXP_REPLACE(column1_utf8, 'e', 'f') AS k
+FROM test;
+----
+Andrfw
+Xiangpfng
+Raphafl
+NULL
+
+# Should run REGEXP_REPLACE with Scalar value for utf8 with flag
+query T
+SELECT
+ REGEXP_REPLACE(column1_utf8, 'e', 'f', 'i') AS k
+FROM test;
+----
+Andrfw
+Xiangpfng
+Raphafl
+NULL
+
+# Should run REGEXP_REPLACE with ScalarArray value for utf8view
+query T
+SELECT
+ REGEXP_REPLACE(column1_utf8view, lower(column1_utf8view), 'bar') AS k
+FROM test;
+----
+Andrew
+Xiangpeng
+Raphael
+NULL
+
+# Should run REGEXP_REPLACE with ScalarArray value for utf8view with flag
+query T
+SELECT
+ REGEXP_REPLACE(column1_utf8view, lower(column1_utf8view), 'bar', 'g') AS k
+FROM test;
+----
+Andrew
+Xiangpeng
+Raphael
+NULL
+
+# Should run REGEXP_REPLACE with ScalarArray value for utf8
+query T
+SELECT
+ REGEXP_REPLACE(column1_utf8, lower(column1_utf8), 'bar') AS k
+FROM test;
+----
+Andrew
+Xiangpeng
+Raphael
+NULL
+
+# Should run REGEXP_REPLACE with ScalarArray value for utf8 with flag
+query T
+SELECT
+ REGEXP_REPLACE(column1_utf8, lower(column1_utf8), 'bar', 'g') AS k
+FROM test;
+----
+Andrew
+Xiangpeng
+Raphael
+NULL
+
### Initcap
query TT
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
