perf: preserve dictionary encoding for bit_length, octet_length, and ascii

datafusion/functions/Cargo.toml

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

datafusion/functions/benches/dictionary_encoding.rs

@@ -0,0 +1,95 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::hint::black_box;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, DictionaryArray};
+use arrow::compute::cast;
+use arrow::datatypes::{Field, Int32Type};
+use criterion::{Criterion, criterion_group, criterion_main};
+use datafusion_common::config::ConfigOptions;
+use datafusion_expr::type_coercion::functions::fields_with_udf;
+use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF};
+
+const NUM_ROWS: usize = 8_192;
+const DICTIONARY_CARDINALITIES: [usize; 4] = [10, 100, 1_000, 8_192];
+
+fn create_string_dictionary(cardinality: usize) -> ArrayRef {
+    let values = (0..NUM_ROWS)
+        .map(|index| Some(format!("value_{:04}", index % cardinality)))
+        .collect::<Vec<_>>();
+    Arc::new(
+        values
+            .iter()
+            .map(|value| value.as_deref())
+            .collect::<DictionaryArray<Int32Type>>(),
+    )
+}
+
+fn benchmark_dictionary_string_udfs(c: &mut Criterion) {
+    let udfs: [(&str, Arc<ScalarUDF>); 3] = [
+        ("ascii", datafusion_functions::string::ascii()),
+        ("bit_length", datafusion_functions::string::bit_length()),
+        ("octet_length", datafusion_functions::string::octet_length()),
+    ];
+    let config_options = Arc::new(ConfigOptions::default());
+
+    for cardinality in DICTIONARY_CARDINALITIES {
+        let dictionary = create_string_dictionary(cardinality);
+        let mut group = c.benchmark_group(format!(
+            "dictionary_encoding/string/cardinality_{cardinality}"
+        ));
+        for (name, udf) in &udfs {
+            let input_field =
+                Field::new("a", dictionary.data_type().clone(), false).into();
+            let coerced_field = fields_with_udf(&[input_field], udf.as_ref())
+                .unwrap()
+                .into_iter()
+                .next()
+                .unwrap();
+            let coerced_type = coerced_field.data_type();
+            let return_type =
+                udf.return_type(std::slice::from_ref(coerced_type)).unwrap();
+            let return_field = Field::new("f", return_type, false).into();
+            let input = if dictionary.data_type() == coerced_type {
+                Arc::clone(&dictionary)
+            } else {
+                cast(dictionary.as_ref(), coerced_type).unwrap()
+            };
+
+            group.bench_function(*name, |b| {
+                b.iter(|| {
+                    black_box(
+                        udf.invoke_with_args(ScalarFunctionArgs {
+                            args: vec![ColumnarValue::Array(Arc::clone(&input))],
+                            arg_fields: vec![Arc::clone(&coerced_field)],
+                            number_rows: NUM_ROWS,
+                            return_field: Arc::clone(&return_field),
+                            config_options: Arc::clone(&config_options),
+                        })
+                        .unwrap(),
+                    )
+                })
+            });
+        }
+        group.finish();
+    }
+}
+
+criterion_group!(benches, benchmark_dictionary_string_udfs);
+criterion_main!(benches);

datafusion/functions/src/string/ascii.rs

@@ -15,13 +15,16 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::utils::transform_leaf_type_preserving_encoding;
 use arrow::array::{ArrayRef, AsArray, Int32Array, StringArrayType};
 use arrow::datatypes::DataType;
 use arrow::error::ArrowError;
 use datafusion_common::types::logical_string;
 use datafusion_common::utils::take_function_args;
 use datafusion_common::{Result, ScalarValue, internal_err};
-use datafusion_expr::{ColumnarValue, Documentation, TypeSignatureClass};
+use datafusion_expr::{
+    ColumnarValue, Documentation, EncodingPreservation, TypeSignatureClass,
+};
 use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
 use datafusion_expr_common::signature::Coercion;
 use datafusion_macros::user_doc;
@@ -63,9 +66,10 @@ impl AsciiFunc {
     pub fn new() -> Self {
         Self {
             signature: Signature::coercible(
-                vec![Coercion::new_exact(TypeSignatureClass::Native(
-                    logical_string(),
-                ))],
+                vec![
+                    Coercion::new_exact(TypeSignatureClass::Native(logical_string()))
+                        .with_encoding_preservation(EncodingPreservation::dictionary()),
+                ],
                 Volatility::Immutable,
             ),
         }
@@ -81,33 +85,16 @@ impl ScalarUDFImpl for AsciiFunc {
         &self.signature
     }
 
-    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
-        Ok(DataType::Int32)
+    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+        transform_leaf_type_preserving_encoding(&arg_types[0], &|_| Ok(DataType::Int32))
     }
 
     fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
         let [arg] = take_function_args(self.name(), args.args)?;
 
         match arg {
             ColumnarValue::Scalar(scalar) => {
-                if scalar.is_null() {
-                    return Ok(ColumnarValue::Scalar(ScalarValue::Int32(None)));
-                }
-
-                match scalar {
-                    ScalarValue::Utf8(Some(s))
-                    | ScalarValue::LargeUtf8(Some(s))
-                    | ScalarValue::Utf8View(Some(s)) => {
-                        let result = first_char_code(&s);
-                        Ok(ColumnarValue::Scalar(ScalarValue::Int32(Some(result))))
-                    }
-                    _ => {
-                        internal_err!(
-                            "Unexpected data type {:?} for function ascii",
-                            scalar.data_type()
-                        )
-                    }
-                }
+                Ok(ColumnarValue::Scalar(ascii_scalar(&scalar)?))
             }
             ColumnarValue::Array(array) => Ok(ColumnarValue::Array(ascii(&[array])?)),
         }
@@ -118,6 +105,24 @@ impl ScalarUDFImpl for AsciiFunc {
     }
 }
 
+fn ascii_scalar(scalar: &ScalarValue) -> Result<ScalarValue> {
+    match scalar {
+        ScalarValue::Utf8(value)
+        | ScalarValue::LargeUtf8(value)
+        | ScalarValue::Utf8View(value) => {
+            Ok(ScalarValue::Int32(value.as_deref().map(first_char_code)))
+        }
+        ScalarValue::Dictionary(key_type, value) => Ok(ScalarValue::Dictionary(
+            key_type.clone(),
+            Box::new(ascii_scalar(value)?),
+        )),
+        _ => internal_err!(
+            "Unexpected data type {:?} for function ascii",
+            scalar.data_type()
+        ),
+    }
+}
+
 /// Returns the Unicode scalar value of the first character of `s`, or 0 when
 /// `s` is empty. Reads the leading byte first so the common all-ASCII case
 /// avoids constructing a `char` iterator and decoding a multi-byte sequence.
@@ -184,6 +189,11 @@ pub fn ascii(args: &[ArrayRef]) -> Result<ArrayRef> {
             let string_array = args[0].as_string_view();
             Ok(calculate_ascii(&string_array)?)
         }
+        DataType::Dictionary(_, _) => {
+            let dictionary = args[0].as_any_dictionary();
+            let converted = ascii(&[Arc::clone(dictionary.values())])?;
+            Ok(dictionary.with_values(converted))
+        }
         _ => internal_err!("Unsupported data type"),
     }
 }

datafusion/functions/src/string/bit_length.rs

@@ -18,13 +18,13 @@
 use arrow::compute::kernels::length::bit_length;
 use arrow::datatypes::DataType;
 
-use crate::utils::utf8_to_int_type;
+use crate::utils::{transform_leaf_type_preserving_encoding, utf8_to_int_type};
 use datafusion_common::types::logical_string;
 use datafusion_common::utils::take_function_args;
 use datafusion_common::{Result, ScalarValue};
 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;
 
@@ -59,9 +59,10 @@ impl BitLengthFunc {
     pub fn new() -> Self {
         Self {
             signature: Signature::coercible(
-                vec![Coercion::new_exact(TypeSignatureClass::Native(
-                    logical_string(),
-                ))],
+                vec![
+                    Coercion::new_exact(TypeSignatureClass::Native(logical_string()))
+                        .with_encoding_preservation(EncodingPreservation::dictionary()),
+                ],
                 Volatility::Immutable,
             ),
         }
@@ -78,30 +79,39 @@ impl ScalarUDFImpl for BitLengthFunc {
     }
 
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
-        utf8_to_int_type(&arg_types[0], "bit_length")
+        transform_leaf_type_preserving_encoding(&arg_types[0], &|data_type| {
+            utf8_to_int_type(data_type, "bit_length")
+        })
     }
 
     fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
         let [array] = take_function_args(self.name(), &args.args)?;
 
         match array {
             ColumnarValue::Array(v) => Ok(ColumnarValue::Array(bit_length(v.as_ref())?)),
-            ColumnarValue::Scalar(v) => match v {
-                ScalarValue::Utf8(v) => Ok(ColumnarValue::Scalar(ScalarValue::Int32(
-                    v.as_ref().map(|x| (x.len() * 8) as i32),
-                ))),
-                ScalarValue::LargeUtf8(v) => Ok(ColumnarValue::Scalar(
-                    ScalarValue::Int64(v.as_ref().map(|x| (x.len() * 8) as i64)),
-                )),
-                ScalarValue::Utf8View(v) => Ok(ColumnarValue::Scalar(
-                    ScalarValue::Int32(v.as_ref().map(|x| (x.len() * 8) as i32)),
-                )),
-                _ => unreachable!("bit length"),
-            },
+            ColumnarValue::Scalar(v) => Ok(ColumnarValue::Scalar(bit_length_scalar(v))),
         }
     }
 
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
 }
+
+fn bit_length_scalar(value: &ScalarValue) -> ScalarValue {
+    match value {
+        ScalarValue::Utf8(v) => {
+            ScalarValue::Int32(v.as_ref().map(|x| (x.len() * 8) as i32))
+        }
+        ScalarValue::LargeUtf8(v) => {
+            ScalarValue::Int64(v.as_ref().map(|x| (x.len() * 8) as i64))
+        }
+        ScalarValue::Utf8View(v) => {
+            ScalarValue::Int32(v.as_ref().map(|x| (x.len() * 8) as i32))
+        }
+        ScalarValue::Dictionary(key_type, value) => {
+            ScalarValue::Dictionary(key_type.clone(), Box::new(bit_length_scalar(value)))
+        }
+        _ => unreachable!("bit length"),
+    }
+}

datafusion/functions/src/string/octet_length.rs

@@ -18,13 +18,13 @@
 use arrow::compute::kernels::length::length;
 use arrow::datatypes::DataType;
 
-use crate::utils::utf8_to_int_type;
+use crate::utils::{transform_leaf_type_preserving_encoding, utf8_to_int_type};
 use datafusion_common::types::logical_string;
 use datafusion_common::utils::take_function_args;
 use datafusion_common::{Result, ScalarValue};
 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;
 
@@ -59,9 +59,10 @@ impl OctetLengthFunc {
     pub fn new() -> Self {
         Self {
             signature: Signature::coercible(
-                vec![Coercion::new_exact(TypeSignatureClass::Native(
-                    logical_string(),
-                ))],
+                vec![
+                    Coercion::new_exact(TypeSignatureClass::Native(logical_string()))
+                        .with_encoding_preservation(EncodingPreservation::dictionary()),
+                ],
                 Volatility::Immutable,
             ),
         }
@@ -78,26 +79,17 @@ impl ScalarUDFImpl for OctetLengthFunc {
     }
 
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
-        utf8_to_int_type(&arg_types[0], "octet_length")
+        transform_leaf_type_preserving_encoding(&arg_types[0], &|data_type| {
+            utf8_to_int_type(data_type, "octet_length")
+        })
     }
 
     fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
         let [array] = take_function_args(self.name(), &args.args)?;
 
         match array {
             ColumnarValue::Array(v) => Ok(ColumnarValue::Array(length(v.as_ref())?)),
-            ColumnarValue::Scalar(v) => match v {
-                ScalarValue::Utf8(v) => Ok(ColumnarValue::Scalar(ScalarValue::Int32(
-                    v.as_ref().map(|x| x.len() as i32),
-                ))),
-                ScalarValue::LargeUtf8(v) => Ok(ColumnarValue::Scalar(
-                    ScalarValue::Int64(v.as_ref().map(|x| x.len() as i64)),
-                )),
-                ScalarValue::Utf8View(v) => Ok(ColumnarValue::Scalar(
-                    ScalarValue::Int32(v.as_ref().map(|x| x.len() as i32)),
-                )),
-                _ => unreachable!("OctetLengthFunc"),
-            },
+            ColumnarValue::Scalar(v) => Ok(ColumnarValue::Scalar(octet_length_scalar(v))),
         }
     }
 
@@ -106,6 +98,23 @@ impl ScalarUDFImpl for OctetLengthFunc {
     }
 }
 
+fn octet_length_scalar(value: &ScalarValue) -> ScalarValue {
+    match value {
+        ScalarValue::Utf8(v) => ScalarValue::Int32(v.as_ref().map(|x| x.len() as i32)),
+        ScalarValue::LargeUtf8(v) => {
+            ScalarValue::Int64(v.as_ref().map(|x| x.len() as i64))
+        }
+        ScalarValue::Utf8View(v) => {
+            ScalarValue::Int32(v.as_ref().map(|x| x.len() as i32))
+        }
+        ScalarValue::Dictionary(key_type, value) => ScalarValue::Dictionary(
+            key_type.clone(),
+            Box::new(octet_length_scalar(value)),
+        ),
+        _ => unreachable!("OctetLengthFunc"),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;