Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Initial support for regex_replace on StringViewArray #11556

Merged
merged 2 commits into from
Jul 22, 2024
Merged

Initial support for regex_replace on StringViewArray #11556

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6606a60
initial support for string view regex
XiangpengHao Jul 19, 2024
2c5cdc8
update tests
XiangpengHao Jul 19, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
initial support for string view regex
  • Loading branch information
@XiangpengHao
XiangpengHao committed Jul 19, 2024
commit 6606a607f6db4aef681f21520b5cb1f90f114fcc
207 changes: 145 additions & 62 deletions datafusion/functions/src/regex/regexpreplace.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,14 @@

//! 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::{Array, ArrayRef, OffsetSizeTrait};
use arrow::datatypes::DataType;
use datafusion_common::cast::as_string_view_array;
use datafusion_common::exec_err;
use datafusion_common::plan_err;
use datafusion_common::ScalarValue;
Expand Down Expand Up @@ -54,6 +57,7 @@ impl RegexpReplaceFunc {
signature: Signature::one_of(
vec![
Exact(vec![Utf8, Utf8, Utf8]),
Exact(vec![Utf8View, Utf8, Utf8]),
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here I added Exact(vec![Utf8View, Utf8, Utf8]), but I wonder if other combinations might also work, e.g., Exact(vec![LargeUtf8, Utf8, Utf8]), Exact(vec![Utf8View, Utf8, LargeUtf8]) etc.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah -- this is the kind of things that the LogicalType proposal #11513 https://github.com/notfilippo would I think make eaiser

Exact(vec![Utf8, Utf8, Utf8, Utf8]),
],
Volatility::Immutable,
Expand All @@ -80,6 +84,7 @@ impl ScalarUDFImpl for RegexpReplaceFunc {
Ok(match &arg_types[0] {
LargeUtf8 | LargeBinary => LargeUtf8,
Utf8 | Binary => Utf8,
Utf8View | BinaryView => Utf8View,
Null => Null,
Dictionary(_, t) => match **t {
LargeUtf8 | LargeBinary => LargeUtf8,
Expand Down Expand Up @@ -118,15 +123,18 @@ impl ScalarUDFImpl for RegexpReplaceFunc {
}
}
}

fn regexp_replace_func(args: &[ColumnarValue]) -> Result<ArrayRef> {
match args[0].data_type() {
DataType::Utf8 => specialize_regexp_replace::<i32>(args),
DataType::LargeUtf8 => specialize_regexp_replace::<i64>(args),
DataType::Utf8View => specialize_regexp_replace::<i32>(args),
other => {
internal_err!("Unsupported data type {other:?} for function regexp_replace")
}
}
}

/// replace POSIX capture groups (like \1) with Rust Regex group (like ${1})
/// used by regexp_replace
fn regex_replace_posix_groups(replacement: &str) -> String {
Expand Down Expand Up @@ -280,8 +288,8 @@ pub fn regexp_replace<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef>
}
}

fn _regexp_replace_early_abort<T: OffsetSizeTrait>(
input_array: &GenericStringArray<T>,
fn _regexp_replace_early_abort<T: ArrayAccessor>(
input_array: T,
sz: usize,
) -> Result<ArrayRef> {
// Mimicking the existing behavior of regexp_replace, if any of the scalar arguments
Expand All @@ -290,6 +298,7 @@ fn _regexp_replace_early_abort<T: OffsetSizeTrait>(
// Also acts like an early abort mechanism when the input array is empty.
Ok(new_null_array(input_array.data_type(), sz))
}

/// Get the first argument from the given string array.
///
/// Note: If the array is empty or the first argument is null,
Expand All @@ -313,8 +322,7 @@ macro_rules! fetch_string_arg {
fn _regexp_replace_static_pattern_replace<T: OffsetSizeTrait>(
args: &[ArrayRef],
) -> Result<ArrayRef> {
let string_array = as_generic_string_array::<T>(&args[0])?;
let array_size = string_array.len();
let array_size = args[0].len();
let pattern = fetch_string_arg!(
&args[1],
"pattern",
Expand Down Expand Up @@ -358,32 +366,61 @@ fn _regexp_replace_static_pattern_replace<T: OffsetSizeTrait>(
// with rust ones.
let replacement = regex_replace_posix_groups(replacement);

// We are going to create the underlying string buffer from its parts
// to be able to re-use the existing null buffer for sparse arrays.
let mut vals = BufferBuilder::<u8>::new({
let offsets = string_array.value_offsets();
(offsets[string_array.len()] - offsets[0])
.to_usize()
.expect("Failed to convert usize")
});
let mut new_offsets = BufferBuilder::<T>::new(string_array.len() + 1);
new_offsets.append(T::zero());

string_array.iter().for_each(|val| {
if let Some(val) = val {
let result = re.replacen(val, limit, replacement.as_str());
vals.append_slice(result.as_bytes());
let string_array_type = args[0].data_type();
match string_array_type {
DataType::Utf8 | DataType::LargeUtf8 => {
let string_array = as_generic_string_array::<T>(&args[0])?;

// We are going to create the underlying string buffer from its parts
// to be able to re-use the existing null buffer for sparse arrays.
let mut vals = BufferBuilder::<u8>::new({
let offsets = string_array.value_offsets();
(offsets[string_array.len()] - offsets[0])
.to_usize()
.unwrap()
});
let mut new_offsets = BufferBuilder::<T>::new(string_array.len() + 1);
new_offsets.append(T::zero());

string_array.iter().for_each(|val| {
if let Some(val) = val {
let result = re.replacen(val, limit, replacement.as_str());
vals.append_slice(result.as_bytes());
}
new_offsets.append(T::from_usize(vals.len()).unwrap());
});

let data = ArrayDataBuilder::new(GenericStringArray::<T>::DATA_TYPE)
.len(string_array.len())
.nulls(string_array.nulls().cloned())
.buffers(vec![new_offsets.finish(), vals.finish()])
.build()?;
let result_array = GenericStringArray::<T>::from(data);
Ok(Arc::new(result_array) as ArrayRef)
}
new_offsets.append(T::from_usize(vals.len()).unwrap());
});

let data = ArrayDataBuilder::new(GenericStringArray::<T>::DATA_TYPE)
.len(string_array.len())
.nulls(string_array.nulls().cloned())
.buffers(vec![new_offsets.finish(), vals.finish()])
.build()?;
let result_array = GenericStringArray::<T>::from(data);
Ok(Arc::new(result_array) as ArrayRef)
DataType::Utf8View => {
let string_view_array = as_string_view_array(&args[0])?;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is the idea that we could specialize this implementation to potentially do inplace updates or something if the string replacements allowed (e.g. replacing larger strings with smaller)?


let mut builder = StringViewBuilder::with_capacity(string_view_array.len())
.with_block_size(1024 * 1024 * 2);

for val in string_view_array.iter() {
if let Some(val) = val {
let result = re.replacen(val, limit, replacement.as_str());
builder.append_value(result);
} else {
builder.append_null();
}
}

let result = builder.finish();
Ok(Arc::new(result) as ArrayRef)
}
_ => unreachable!(
"Invalid data type for regexp_replace: {}",
string_array_type
),
}
}

/// Determine which implementation of the regexp_replace to use based
Expand Down Expand Up @@ -469,43 +506,89 @@ mod tests {

use super::*;

#[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);
macro_rules! static_pattern_regexp_replace {
($name:ident, $T:ty) => {
#[test]
fn $name() {
let values = vec!["abc", "acd", "abcd1234567890123", "123456789012abc"];
let patterns = vec!["b"; 4];
let replacement = vec!["foo"; 4];
let expected =
vec!["afooc", "acd", "afoocd1234567890123", "123456789012afooc"];

let values = <$T>::from(values);
let patterns = StringArray::from(patterns);
let replacements = StringArray::from(replacement);
let expected = <$T>::from(expected);

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);
static_pattern_regexp_replace!(string_array, StringArray);
static_pattern_regexp_replace!(string_view_array, StringViewArray);
static_pattern_regexp_replace!(large_string_array, LargeStringArray);

macro_rules! static_pattern_regexp_replace_with_flags {
($name:ident, $T:ty) => {
#[test]
fn $name() {
let values = vec![
"abc",
"aBc",
"acd",
"abcd1234567890123",
"aBcd1234567890123",
"123456789012abc",
"123456789012aBc",
];
let expected = vec![
"afooc",
"afooc",
"acd",
"afoocd1234567890123",
"afoocd1234567890123",
"123456789012afooc",
"123456789012afooc",
];

let values = <$T>::from(values);
let patterns = StringArray::from(vec!["b"; 7]);
let replacements = StringArray::from(vec!["foo"; 7]);
let flags = StringArray::from(vec!["i"; 5]);
let expected = <$T>::from(expected);

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);
}
};
}

static_pattern_regexp_replace_with_flags!(string_array_with_flags, StringArray);
static_pattern_regexp_replace_with_flags!(
string_view_array_with_flags,
StringViewArray
);
static_pattern_regexp_replace_with_flags!(
large_string_array_with_flags,
LargeStringArray
);

#[test]
fn test_static_pattern_regexp_replace_early_abort() {
let values = StringArray::from(vec!["abc"; 5]);
Expand Down
Loading