diff --git a/Cargo.lock b/Cargo.lock index 174010138c3c2..88b400c26ea4d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2274,6 +2274,7 @@ version = "55.1.0" dependencies = [ "arrow", "arrow-buffer", + "arrow-select", "base64 0.23.1", "blake2", "blake3", diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index 28765bdda9e2e..e3906760f7562 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -67,6 +67,7 @@ name = "datafusion_functions" [dependencies] arrow = { workspace = true } arrow-buffer = { workspace = true } +arrow-select = { workspace = true } base64 = { version = "0.23", optional = true } blake2 = { version = "^0.11.0", optional = true } blake3 = { version = "1.8", optional = true } diff --git a/datafusion/functions/src/regex/regexpreplace.rs b/datafusion/functions/src/regex/regexpreplace.rs index 647abfed59086..31e90c636c4ef 100644 --- a/datafusion/functions/src/regex/regexpreplace.rs +++ b/datafusion/functions/src/regex/regexpreplace.rs @@ -18,6 +18,7 @@ //! Regex expressions use memchr::memchr; +use crate::utils::transform_leaf_type_preserving_encoding; use arrow::array::ArrayDataBuilder; use arrow::array::BufferBuilder; use arrow::array::GenericStringArray; @@ -25,7 +26,8 @@ use arrow::array::StringViewBuilder; use arrow::array::{Array, ArrayRef, OffsetSizeTrait}; use arrow::array::{ArrayAccessor, StringViewArray}; use arrow::array::{ArrayIter, AsArray, new_null_array}; -use arrow::datatypes::DataType; +use arrow::datatypes::{DataType, Field, FieldRef}; +use arrow_select::dictionary::garbage_collect_any_dictionary; use datafusion_common::ScalarValue; use datafusion_common::cast::{ as_large_string_array, as_string_array, as_string_view_array, @@ -35,10 +37,10 @@ use datafusion_common::plan_err; use datafusion_common::utils::offset_span_len; use datafusion_common::{Result, cast::as_generic_string_array, internal_err}; use datafusion_expr::ColumnarValue; -use datafusion_expr::TypeSignature; use datafusion_expr::function::Hint; use datafusion_expr::{ - Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, + Documentation, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, }; use datafusion_macros::user_doc; use regex::{CaptureLocations, Regex}; @@ -95,16 +97,8 @@ impl Default for RegexpReplaceFunc { impl RegexpReplaceFunc { pub fn new() -> Self { - use DataType::*; - use TypeSignature::*; Self { - signature: Signature::one_of( - vec![ - Uniform(3, vec![Utf8View, LargeUtf8, Utf8]), - Uniform(4, vec![Utf8View, LargeUtf8, Utf8]), - ], - Volatility::Immutable, - ), + signature: Signature::user_defined(Volatility::Immutable), } } } @@ -119,33 +113,96 @@ impl ScalarUDFImpl for RegexpReplaceFunc { } fn return_type(&self, arg_types: &[DataType]) -> Result { - use DataType::*; - Ok(match &arg_types[0] { - LargeUtf8 | LargeBinary => LargeUtf8, - Utf8 | Binary => Utf8, - Utf8View | BinaryView => Utf8View, - Null => Null, - Dictionary(_, t) => match **t { - LargeUtf8 | LargeBinary => LargeUtf8, - Utf8 | Binary => Utf8, - Null => Null, - _ => { - return plan_err!( - "the regexp_replace can only accept strings but got {:?}", - **t - ); - } - }, - other => { - return plan_err!( - "The regexp_replace function can only accept strings. Got {other}" - ); + transform_leaf_type_preserving_encoding(&arg_types[0], ®exp_replace_leaf_type) + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + let input = &args.arg_fields[0]; + let others_scalar = args.scalar_arguments.iter().skip(1).all(Option::is_some); + let output_type = match input.data_type() { + DataType::Dictionary(_, value_type) if !others_scalar => { + regexp_replace_leaf_type(value_type.as_ref())? } - }) + other => self.return_type(std::slice::from_ref(other))?, + }; + Ok(Arc::new(Field::new(self.name(), output_type, true))) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + if !(3..=4).contains(&arg_types.len()) { + return plan_err!( + "regexp_replace expects 3 or 4 arguments, got {}", + arg_types.len() + ); + } + let leaf = match &arg_types[0] { + DataType::Dictionary(_, v) => v.as_ref(), + other => other, + }; + let is_string_leaf = matches!( + leaf, + DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Utf8View + | DataType::Binary + | DataType::LargeBinary + | DataType::BinaryView + ); + let target = if is_string_leaf { + regexp_replace_leaf_type(leaf)? + } else { + DataType::Utf8View + }; + let first = match &arg_types[0] { + DataType::Dictionary(k, _) if is_string_leaf => { + DataType::Dictionary(k.clone(), Box::new(target.clone())) + } + _ => target.clone(), + }; + let mut out = Vec::with_capacity(arg_types.len()); + out.push(first); + for _ in 1..arg_types.len() { + out.push(target.clone()); + } + Ok(out) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - let args = &args.args; + let preserve_dictionary = + matches!(args.return_field.data_type(), DataType::Dictionary(_, _)); + let mut args = args.args; + + let scalar_dict_key = if let Some(ColumnarValue::Scalar( + ScalarValue::Dictionary(key_type, value), + )) = args.first() + { + let key = key_type.clone(); + args[0] = ColumnarValue::Scalar(value.as_ref().clone()); + preserve_dictionary.then_some(key) + } else { + None + }; + + if let ColumnarValue::Array(array) = &args[0] + && let DataType::Dictionary(_, value_type) = array.data_type().clone() + { + if preserve_dictionary { + let compacted; + let mut dict = array.as_any_dictionary(); + if dict.values().len() > dict.keys().len() { + // Filtered batches and shared Parquet dictionaries can have more + // values than rows, so drop unreferenced values first. + compacted = garbage_collect_any_dictionary(dict)?; + dict = compacted.as_any_dictionary(); + } + let mut inner: Vec = Vec::with_capacity(args.len()); + inner.push(ColumnarValue::Array(Arc::clone(dict.values()))); + inner.extend_from_slice(&args[1..]); + let values_out = regexp_replace_func(&inner)?; + return Ok(ColumnarValue::Array(dict.with_values(values_out))); + } + args[0] = ColumnarValue::Array(arrow::compute::cast(array, &value_type)?); + } let len = args .iter() @@ -155,14 +212,20 @@ impl ScalarUDFImpl for RegexpReplaceFunc { }); let is_scalar = len.is_none(); - let result = regexp_replace_func(args); - if is_scalar { - // If all inputs are scalar, keeps output as scalar - let result = result.and_then(|arr| ScalarValue::try_from_array(&arr, 0)); - result.map(ColumnarValue::Scalar) + let result = regexp_replace_func(&args)?; + let out = if is_scalar { + ColumnarValue::Scalar(ScalarValue::try_from_array(&result, 0)?) } else { - result.map(ColumnarValue::Array) + ColumnarValue::Array(result) + }; + + if let (Some(key_type), ColumnarValue::Scalar(scalar)) = (scalar_dict_key, &out) { + return Ok(ColumnarValue::Scalar(ScalarValue::Dictionary( + key_type, + Box::new(scalar.clone()), + ))); } + Ok(out) } fn documentation(&self) -> Option<&Documentation> { @@ -170,6 +233,21 @@ impl ScalarUDFImpl for RegexpReplaceFunc { } } +fn regexp_replace_leaf_type(dt: &DataType) -> Result { + use DataType::*; + Ok(match dt { + LargeUtf8 | LargeBinary => LargeUtf8, + Utf8 | Binary => Utf8, + Utf8View | BinaryView => Utf8View, + Null => Null, + other => { + return plan_err!( + "the regexp_replace can only accept strings but got {other:?}" + ); + } + }) +} + fn regexp_replace_func(args: &[ColumnarValue]) -> Result { match args[0].data_type() { DataType::Utf8 => specialize_regexp_replace::(args), diff --git a/datafusion/sqllogictest/test_files/regexp/regexp_replace.slt b/datafusion/sqllogictest/test_files/regexp/regexp_replace.slt index 937c41dcf5d13..59c7c7bdf04c5 100644 --- a/datafusion/sqllogictest/test_files/regexp/regexp_replace.slt +++ b/datafusion/sqllogictest/test_files/regexp/regexp_replace.slt @@ -138,6 +138,89 @@ SELECT regexp_replace(arrow_cast('foobar', 'Dictionary(Int32, Utf8)'), 'bar', 'x ---- fooxx +statement ok +CREATE TABLE regexp_replace_dict AS +SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS url +FROM (VALUES + ('https://www.example.com/path'), + ('http://test.org/foo/bar'), + ('https://example.com/') +) AS t(column1); + +query T +SELECT DISTINCT arrow_typeof(regexp_replace(url, '^https?://(?:www\.)?([^/]+)/.*$', '\1')) +FROM regexp_replace_dict; +---- +Dictionary(Int32, Utf8) + +statement ok +DROP TABLE regexp_replace_dict; + +statement ok +CREATE TABLE regexp_replace_dict_runtime_scalars AS +SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS s, + column2 AS p +FROM (VALUES ('abc', 'a'), ('abd', 'b')) AS v; + +query T +SELECT regexp_replace(s, 'a', arrow_typeof(p)) +FROM regexp_replace_dict_runtime_scalars +ORDER BY p; +---- +Utf8bc +Utf8bd + +statement ok +DROP TABLE regexp_replace_dict_runtime_scalars; + +statement ok +CREATE TABLE regexp_replace_dict_types AS SELECT + arrow_cast(column1, 'Dictionary(Int32, LargeUtf8)') AS dl, + arrow_cast(column1, 'Dictionary(Int32, Utf8View)') AS dv, + arrow_cast(column1, 'Dictionary(Int32, Binary)') AS db, + arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS d8, + column1 AS s +FROM (VALUES ('abc'), ('abd'), (NULL)) AS v; + +query TT +SELECT arrow_typeof(regexp_replace(dl, 'a', 'x')), regexp_replace(dl, 'a', 'x') +FROM regexp_replace_dict_types +ORDER BY s NULLS LAST; +---- +Dictionary(Int32, LargeUtf8) xbc +Dictionary(Int32, LargeUtf8) xbd +Dictionary(Int32, LargeUtf8) NULL + +query TT +SELECT arrow_typeof(regexp_replace(dv, 'a', 'x')), regexp_replace(dv, 'a', 'x') +FROM regexp_replace_dict_types +ORDER BY s NULLS LAST; +---- +Dictionary(Int32, Utf8View) xbc +Dictionary(Int32, Utf8View) xbd +Dictionary(Int32, Utf8View) NULL + +query TT +SELECT arrow_typeof(regexp_replace(db, 'a', 'x')), regexp_replace(db, 'a', 'x') +FROM regexp_replace_dict_types +ORDER BY s NULLS LAST; +---- +Dictionary(Int32, Utf8) xbc +Dictionary(Int32, Utf8) xbd +Dictionary(Int32, Utf8) NULL + +query TT +SELECT arrow_typeof(regexp_replace(d8, s, 'x')), regexp_replace(d8, s, 'x') +FROM regexp_replace_dict_types +ORDER BY s NULLS LAST; +---- +Utf8 x +Utf8 x +Utf8 NULL + +statement ok +DROP TABLE regexp_replace_dict_types; + query TTT select regexp_replace(col, NULL, 'c'),