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`.
This is the same cost issue as the `substr` fast path in #25874, and #26696
discusses handling it once for all functions that map over dictionary values,
so this guard may later move into a shared helper.
--
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]
