perf: preserve dictionary encoding through regexp_replace

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 }

datafusion/functions/src/regex/regexpreplace.rs

@@ -18,14 +18,16 @@
 //! 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 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<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 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<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()
@@ -155,21 +212,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),

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'),