jayzhan211 commented on code in PR #26186:
URL: https://github.com/apache/datafusion/pull/26186#discussion_r4366296564
##########
datafusion/functions/src/regex/regexpreplace.rs:
##########
@@ -119,33 +112,89 @@ impl ScalarUDFImpl for RegexpReplaceFunc {
}
fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
- 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<FieldRef> {
+ 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<Vec<DataType>> {
+ 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<ColumnarValue> {
- 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 dict = array.as_any_dictionary();
Review Comment:
This runs the regex over every value in the dictionary, including values no
row references. After a filter, or with a Parquet dictionary shared across
batches, that can be much larger than the batch. In a same-build microbench
(URL-shaped values, the Q28 pattern) I measured 5.3x slower than the old cast
path at 50k values / 8192 rows, and 34x slower at 50k values / 1000 rows.
Compacting the dictionary first fixes it. With this change every case I
tried was ≤1.01x the old path, and the small-dictionary win stays:
```diff
if preserve_dictionary {
- let dict = array.as_any_dictionary();
+ // A filtered batch or a shared Parquet dictionary can hold
far more
+ // values than rows; drop unreferenced values so the regex
runs at
+ // most once per row.
+ let compacted;
+ let mut dict = array.as_any_dictionary();
+ if dict.values().len() > dict.keys().len() {
+ compacted = garbage_collect_any_dictionary(dict)?;
+ dict = compacted.as_any_dictionary();
+ }
```
with `use arrow_select::dictionary::garbage_collect_any_dictionary;` and
`arrow-select = { workspace = true }` added to
`datafusion/functions/Cargo.toml`.
##########
datafusion/sqllogictest/test_files/regexp/regexp_replace.slt:
##########
@@ -138,6 +138,41 @@ 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;
Review Comment:
Could you add a few cases for the other dictionary value types and the
non-literal fallback? `Dictionary(Int32, Utf8View)` used to be a plan error and
now works, and `Binary` values now come back as `Dictionary(Int32, Utf8)`. I
checked these locally and they all return the right values. Fine to do in a
follow-up if you prefer.
```sql
CREATE TABLE t 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;
-- expect Dictionary(Int32, LargeUtf8) / Dictionary(Int32, Utf8View) /
Dictionary(Int32, Utf8)
SELECT arrow_typeof(regexp_replace(dl, 'a', 'x')), regexp_replace(dl, 'a',
'x') FROM t;
SELECT arrow_typeof(regexp_replace(dv, 'a', 'x')), regexp_replace(dv, 'a',
'x') FROM t;
SELECT arrow_typeof(regexp_replace(db, 'a', 'x')), regexp_replace(db, 'a',
'x') FROM t;
-- non-literal pattern: expect Utf8
SELECT arrow_typeof(regexp_replace(d8, s, 'x')), regexp_replace(d8, s, 'x')
FROM t;
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
