From c11a73cd55afe876a4cbc82381ffeec21d4ee6e6 Mon Sep 17 00:00:00 2001 From: RIchard Baah Date: Fri, 9 Oct 2026 16:21:02 -0400 Subject: [PATCH] perf: preserve dictionary encoding in substr For scalar start/length, substr can run over the dict's unique values and reuse the keys instead of casting the dict to its value type. - Signature uses EncodingPreservation::dictionary() so coercion keeps the dict wrapper. - Fast path in invoke_with_args runs substr on dict.values() and rewraps via with_values(..). - Per-row start/length still materializes the dict. - return_field_from_args declares Dictionary for the scalar case and the value type for the per-row case so the planner's expected schema matches runtime output. Closes #26170 --- datafusion/functions/src/unicode/substr.rs | 99 +++++++++++++++++-- .../test_files/string/dictionary_utf8.slt | 6 ++ 2 files changed, 97 insertions(+), 8 deletions(-) diff --git a/datafusion/functions/src/unicode/substr.rs b/datafusion/functions/src/unicode/substr.rs index 55da0e801a6ad..b39343eccbb69 100644 --- a/datafusion/functions/src/unicode/substr.rs +++ b/datafusion/functions/src/unicode/substr.rs @@ -24,15 +24,16 @@ use arrow::array::{ StringArrayType, StringViewArray, }; use arrow::buffer::{NullBuffer, ScalarBuffer}; -use arrow::datatypes::DataType; +use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::cast::as_int64_array; use datafusion_common::types::{ NativeType, logical_int32, logical_int64, logical_string, }; use datafusion_common::{Result, exec_err}; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignature, TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ReturnFieldArgs, + ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, TypeSignatureClass, + Volatility, }; use datafusion_macros::user_doc; @@ -73,7 +74,8 @@ impl Default for SubstrFunc { impl SubstrFunc { pub fn new() -> Self { - let string = Coercion::new_exact(TypeSignatureClass::Native(logical_string())); + let string = Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation(EncodingPreservation::dictionary()); let int64 = Coercion::new_implicit( TypeSignatureClass::Native(logical_int64()), vec![TypeSignatureClass::Native(logical_int32())], @@ -115,8 +117,51 @@ impl ScalarUDFImpl for SubstrFunc { Ok(arg_types[0].clone()) } + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + // Per-row start/length materializes the dict, so the output type is the + // value type; scalar tail keeps the dict wrapper. + let input = &args.arg_fields[0]; + let output_type = match input.data_type() { + DataType::Dictionary(_, value_type) + if !args.scalar_arguments.iter().skip(1).all(Option::is_some) => + { + value_type.as_ref().clone() + } + other => other.clone(), + }; + Ok(Field::new(self.name(), output_type, input.is_nullable()).into()) + } + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(substr, vec![])(&args.args) + if let ColumnarValue::Array(array) = &args.args[0] + && matches!(array.data_type(), DataType::Dictionary(_, _)) + && args.args[1..] + .iter() + .all(|v| matches!(v, ColumnarValue::Scalar(_))) + { + // Fast path: substr the N_distinct values, reuse the keys. + let dict = array.as_any_dictionary(); + let values = Arc::clone(dict.values()); + let mut inner = Vec::with_capacity(args.args.len()); + inner.push(ColumnarValue::Array(values)); + inner.extend_from_slice(&args.args[1..]); + let new_values = make_scalar_function(substr, vec![])(&inner)? + .into_array(dict.values().len())?; + return Ok(ColumnarValue::Array(dict.with_values(new_values))); + } + + // Per-row start/length: materialize the dict so substr can run row-wise. + let args = if let ColumnarValue::Array(array) = &args.args[0] + && let DataType::Dictionary(_, value_type) = array.data_type() + { + let flat = arrow::compute::kernels::cast::cast(array, value_type.as_ref())?; + let mut new_args = args.args.clone(); + new_args[0] = ColumnarValue::Array(flat); + new_args + } else { + args.args + }; + make_scalar_function(substr, vec![])(&args) } fn aliases(&self) -> &[String] { @@ -354,13 +399,15 @@ mod tests { use std::sync::Arc; use arrow::array::{ - Array, ArrayRef, AsArray, Int64Array, LargeStringArray, StringArray, - StringViewArray, + Array, ArrayRef, AsArray, DictionaryArray, Int64Array, LargeStringArray, + StringArray, StringViewArray, }; use arrow::datatypes::DataType::{LargeUtf8, Utf8, Utf8View}; + use arrow::datatypes::{DataType, Field, Int32Type}; + use datafusion_common::config::ConfigOptions; use datafusion_common::{Result, ScalarValue, exec_err}; - use datafusion_expr::{ColumnarValue, ScalarUDFImpl}; + use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; use crate::unicode::substr::SubstrFunc; use crate::utils::test::test_function; @@ -766,4 +813,40 @@ mod tests { Ok(()) } + + #[test] + fn test_substr_dict_fast_path_preserves_encoding() -> Result<()> { + let dict_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(Utf8)); + let keys = arrow::array::Int32Array::from(vec![0, 1, 0, 2, 1]); + let values = StringArray::from(vec!["silver", "copper", "golden"]); + let dict: ArrayRef = Arc::new( + DictionaryArray::::try_new(keys, Arc::new(values)).unwrap(), + ); + + let result = SubstrFunc::new() + .invoke_with_args(ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(dict), + ColumnarValue::Scalar(ScalarValue::from(1i64)), + ColumnarValue::Scalar(ScalarValue::from(3i64)), + ], + arg_fields: vec![ + Arc::new(Field::new("s", dict_type.clone(), true)), + Arc::new(Field::new("start", DataType::Int64, false)), + Arc::new(Field::new("len", DataType::Int64, false)), + ], + number_rows: 5, + return_field: Arc::new(Field::new("r", dict_type, true)), + config_options: Arc::new(ConfigOptions::default()), + })? + .into_array(5)?; + + assert!(matches!(result.data_type(), DataType::Dictionary(_, v) if **v == Utf8)); + let out = result.as_dictionary::(); + let v = out.values().as_string::(); + assert_eq!(v.len(), 3); + assert_eq!((v.value(0), v.value(1), v.value(2)), ("sil", "cop", "gol")); + + Ok(()) + } } diff --git a/datafusion/sqllogictest/test_files/string/dictionary_utf8.slt b/datafusion/sqllogictest/test_files/string/dictionary_utf8.slt index 99271a28f9509..91a12ca52c422 100644 --- a/datafusion/sqllogictest/test_files/string/dictionary_utf8.slt +++ b/datafusion/sqllogictest/test_files/string/dictionary_utf8.slt @@ -57,6 +57,12 @@ SELECT arrow_cast('', 'Dictionary(Int32, Utf8)'); # include ./string_query.slt.part +# Scalar start/length preserves dict encoding end-to-end. +query T +SELECT DISTINCT arrow_typeof(substr(c1, 1, 3)) FROM test_substr; +---- +Dictionary(Int32, Utf8) + # # Clean up #