perf: evaluate elementwise scalar functions once per distinct dictionary value

datafusion/expr/src/udf.rs

@@ -317,6 +317,14 @@ impl ScalarUDF {
         self.inner.short_circuits()
     }
 
+    /// Returns true if each output row of this function depends only on the
+    /// corresponding input row.
+    ///
+    /// See [ScalarUDFImpl::evaluates_elementwise] for more information.
+    pub fn evaluates_elementwise(&self) -> bool {
+        self.inner.evaluates_elementwise()
+    }
+
     /// Computes the output interval for a [`ScalarUDF`], given the input
     /// intervals.
     ///
@@ -885,6 +893,22 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any {
         false
     }
 
+    /// Returns true if each output row depends only on the corresponding input
+    /// row, with no cross-row state.
+    ///
+    /// When true, a dictionary-encoded argument is unwrapped before the call:
+    /// the function is evaluated over the distinct values and the result
+    /// re-mapped through the keys, or, where that does not pay off, over the
+    /// expanded column. [`ScalarFunctionArgs::number_rows`] and the argument
+    /// fields then describe the array actually passed, not the planned batch.
+    ///
+    /// Values that no key references may still be evaluated, so a function that
+    /// can error on valid input should not opt in. Volatile functions are never
+    /// unwrapped, nor are arguments whose field carries metadata.
+    fn evaluates_elementwise(&self) -> bool {
+        false
+    }
+
     /// Determines which of the arguments passed to this function are evaluated eagerly
     /// and which may be evaluated lazily.
     ///
@@ -1182,6 +1206,10 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl {
         self.inner.short_circuits()
     }
 
+    fn evaluates_elementwise(&self) -> bool {
+        self.inner.evaluates_elementwise()
+    }
+
     fn evaluate_bounds(&self, input: &[&Interval]) -> Result<Interval> {
         self.inner.evaluate_bounds(input)
     }

datafusion/functions/Cargo.toml

@@ -403,8 +403,3 @@ required-features = ["math_expressions"]
 harness = false
 name = "round"
 required-features = ["math_expressions"]
-
-[[bench]]
-harness = false
-name = "dictionary_encoding"
-required-features = ["string_expressions", "unicode_expressions"]

datafusion/functions/src/encoding/inner.rs

@@ -39,8 +39,8 @@ use datafusion_common::{
     },
 };
 use datafusion_expr::{
-    Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
-    TypeSignatureClass, Volatility,
+    Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs,
+    ScalarUDFImpl, Signature, TypeSignatureClass, Volatility,
 };
 use datafusion_macros::user_doc;
 use std::fmt;
@@ -94,7 +94,8 @@ impl EncodeFunc {
                         TypeSignatureClass::Binary,
                         vec![TypeSignatureClass::Native(logical_string())],
                         NativeType::Binary,
-                    ),
+                    )
+                    .with_encoding_preservation(EncodingPreservation::dictionary()),
                     Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
                 ],
                 Volatility::Immutable,
@@ -119,6 +120,10 @@ impl ScalarUDFImpl for EncodeFunc {
         }
     }
 
+    fn evaluates_elementwise(&self) -> bool {
+        true
+    }
+
     fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
         let [expression, encoding] = take_function_args("encode", &args.args)?;
         let encoding = Encoding::try_from(encoding)?;

datafusion/functions/src/unicode/initcap.rs

@@ -86,6 +86,10 @@ impl ScalarUDFImpl for InitcapFunc {
         Ok(arg_types[0].clone())
     }
 
+    fn evaluates_elementwise(&self) -> bool {
+        true
+    }
+
     fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
         match &args.args[0] {
             ColumnarValue::Scalar(scalar) => {
@@ -140,6 +144,8 @@ fn initcap_array(array: &ArrayRef) -> Result<ArrayRef> {
         DataType::Utf8 => initcap::<i32>(&[Arc::clone(array)]),
         DataType::LargeUtf8 => initcap::<i64>(&[Arc::clone(array)]),
         DataType::Utf8View => initcap_utf8view(&[Arc::clone(array)]),
+        // Serves the calls the physical layer leaves encoded, such as fields
+        // carrying extension metadata.
         DataType::Dictionary(_, _) => {
             let dictionary = array.as_any_dictionary();
             let converted = initcap_array(dictionary.values())?;

datafusion/functions/src/unicode/reverse.rs

@@ -87,6 +87,10 @@ impl ScalarUDFImpl for ReverseFunc {
         Ok(arg_types[0].clone())
     }
 
+    fn evaluates_elementwise(&self) -> bool {
+        true
+    }
+
     fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
         make_scalar_function(reverse, vec![])(&args.args)
     }
@@ -114,6 +118,8 @@ fn reverse(args: &[ArrayRef]) -> Result<ArrayRef> {
             &args[0].as_string_view(),
             StringViewArrayBuilder::with_capacity(len),
         ),
+        // Serves the calls the physical layer leaves encoded, such as fields
+        // carrying extension metadata.
         DataType::Dictionary(_, _) => {
             let dictionary = args[0].as_any_dictionary();
             let converted = reverse(&[Arc::clone(dictionary.values())])?;

datafusion/physical-expr/Cargo.toml

@@ -102,5 +102,9 @@ name = "simplify"
 harness = false
 name = "string_concat"
 
+[[bench]]
+harness = false
+name = "dictionary_encoding"
+
 [package.metadata.cargo-machete]
 ignored = ["half"]