Skip to content

perf: preserve dictionary encoding through regexp_replace #26186

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

Draft
Rich-T-kid wants to merge 2 commits into
base: main
Choose a base branch
from
+146 −41
Draft

perf: preserve dictionary encoding through regexp_replace #26186

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7a2ee99
perf: preserve dictionary encoding through regexp_replace
Rich-T-kid Oct 12, 2026
80bd16e
fix: honor planned regexp_replace return type
Rich-T-kid Oct 18, 2026
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
152 changes: 111 additions & 41 deletions datafusion/functions/src/regex/regexpreplace.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,14 +18,15 @@
//! 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;
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 datafusion_common::ScalarValue;
use datafusion_common::cast::{
as_large_string_array, as_string_array, as_string_view_array,
Expand All @@ -35,10 +36,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};
Expand Down Expand Up @@ -95,16 +96,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),
}
}
}
Expand All @@ -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], &regexp_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();

@jayzhan211 jayzhan211 Oct 24, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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:

             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.

let mut inner: Vec<ColumnarValue> = 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()
Expand All @@ -155,21 +204,42 @@ 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> {
self.doc()
}
}

fn regexp_replace_leaf_type(dt: &DataType) -> Result<DataType> {
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<ArrayRef> {
match args[0].data_type() {
DataType::Utf8 => specialize_regexp_replace::<i32>(args),
Expand Down
35 changes: 35 additions & 0 deletions datafusion/sqllogictest/test_files/regexp/regexp_replace.slt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

@jayzhan211 jayzhan211 Oct 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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.

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;


query TTT
select
regexp_replace(col, NULL, 'c'),
Expand Down
Loading