From 1896e066fc9ade37dc2fc049e47689305c64e3cf Mon Sep 17 00:00:00 2001 From: lyne7-sc <734432041@qq.com> Date: Thu, 11 Jun 2026 00:31:55 +0800 Subject: [PATCH 1/7] perf: preserve dictionary encoding for case conversion --- datafusion/expr-common/src/signature.rs | 72 +++++++++++- datafusion/expr/src/lib.rs | 4 +- .../expr/src/type_coercion/functions.rs | 104 ++++++++++++++++-- datafusion/functions/src/string/common.rs | 39 ++++++- datafusion/functions/src/string/lower.rs | 65 +++++++++-- datafusion/functions/src/string/upper.rs | 65 +++++++++-- .../sqllogictest/test_files/functions.slt | 4 +- 7 files changed, 323 insertions(+), 30 deletions(-) diff --git a/datafusion/expr-common/src/signature.rs b/datafusion/expr-common/src/signature.rs index 35a679cc447cf..b81290bb19fc5 100644 --- a/datafusion/expr-common/src/signature.rs +++ b/datafusion/expr-common/src/signature.rs @@ -1043,12 +1043,22 @@ fn get_data_types(native_type: &NativeType) -> Vec { /// /// * `Exact` - Only accepts arguments that exactly match the desired type /// * `Implicit` - Accepts the desired type and can coerce from specified source types +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Hash)] +pub enum EncodingPreservation { + /// Do not request preservation of a physical encoding. + None, + /// Preserve dictionary encoding and coerce only the dictionary values. + Dictionary, +} + #[derive(Debug, Clone, Eq, PartialOrd)] pub enum Coercion { /// Coercion that only accepts arguments exactly matching the desired type. Exact { /// The required type for the argument desired_type: TypeSignatureClass, + /// Physical encoding preservation requested by the function. + encoding_preservation: EncodingPreservation, }, /// Coercion that accepts the desired type and can implicitly coerce from other types. @@ -1057,12 +1067,27 @@ pub enum Coercion { desired_type: TypeSignatureClass, /// Rules for implicit coercion from other types implicit_coercion: ImplicitCoercion, + /// Physical encoding preservation requested by the function. + encoding_preservation: EncodingPreservation, }, } impl Coercion { pub fn new_exact(desired_type: TypeSignatureClass) -> Self { - Self::Exact { desired_type } + Self::Exact { + desired_type, + encoding_preservation: EncodingPreservation::None, + } + } + + pub fn new_exact_preserving_encoding( + desired_type: TypeSignatureClass, + encoding_preservation: EncodingPreservation, + ) -> Self { + Self::Exact { + desired_type, + encoding_preservation, + } } /// Create a new coercion with implicit coercion rules. @@ -1080,6 +1105,37 @@ impl Coercion { allowed_source_types, default_casted_type, }, + encoding_preservation: EncodingPreservation::None, + } + } + + pub fn with_encoding_preservation( + mut self, + encoding_preservation: EncodingPreservation, + ) -> Self { + match &mut self { + Coercion::Exact { + encoding_preservation: current, + .. + } + | Coercion::Implicit { + encoding_preservation: current, + .. + } => *current = encoding_preservation, + } + self + } + + pub fn encoding_preservation(&self) -> EncodingPreservation { + match self { + Coercion::Exact { + encoding_preservation, + .. + } + | Coercion::Implicit { + encoding_preservation, + .. + } => *encoding_preservation, } } @@ -1103,7 +1159,7 @@ impl Coercion { pub fn desired_type(&self) -> &TypeSignatureClass { match self { - Coercion::Exact { desired_type } => desired_type, + Coercion::Exact { desired_type, .. } => desired_type, Coercion::Implicit { desired_type, .. } => desired_type, } } @@ -1128,6 +1184,7 @@ impl PartialEq for Coercion { fn eq(&self, other: &Self) -> bool { self.desired_type() == other.desired_type() && self.implicit_coercion() == other.implicit_coercion() + && self.encoding_preservation() == other.encoding_preservation() } } @@ -1135,6 +1192,7 @@ impl Hash for Coercion { fn hash(&self, state: &mut H) { self.desired_type().hash(state); self.implicit_coercion().hash(state); + self.encoding_preservation().hash(state); } } @@ -2178,6 +2236,16 @@ mod tests { assert_snapshot!(implicit_with_multiple_sources, @"Int64"); } + #[test] + fn test_coercion_encoding_preservation_affects_equality() { + let default = Coercion::new_exact(TypeSignatureClass::Native(logical_string())); + let preserving = default + .clone() + .with_encoding_preservation(EncodingPreservation::Dictionary); + + assert_ne!(default, preserving); + } + #[test] fn test_to_string_repr_coercible() { use insta::assert_snapshot; diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs index b52a784df931a..43cb3fdc20c40 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -101,8 +101,8 @@ pub use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator}; pub use datafusion_expr_common::operator::Operator; pub use datafusion_expr_common::placement::ExpressionPlacement; pub use datafusion_expr_common::signature::{ - ArrayFunctionArgument, ArrayFunctionSignature, Coercion, Signature, - TIMEZONE_WILDCARD, TypeSignature, TypeSignatureClass, Volatility, + ArrayFunctionArgument, ArrayFunctionSignature, Coercion, EncodingPreservation, + Signature, TIMEZONE_WILDCARD, TypeSignature, TypeSignatureClass, Volatility, }; pub use datafusion_expr_common::type_coercion::binary; pub use expr::{ diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs index 33746a2c46b30..640de14e40bfd 100644 --- a/datafusion/expr/src/type_coercion/functions.rs +++ b/datafusion/expr/src/type_coercion/functions.rs @@ -33,7 +33,7 @@ use datafusion_common::utils::{ use datafusion_common::{ Result, exec_err, internal_err, plan_err, types::NativeType, utils::list_ndims, }; -use datafusion_expr_common::signature::ArrayFunctionArgument; +use datafusion_expr_common::signature::{ArrayFunctionArgument, EncodingPreservation}; use datafusion_expr_common::type_coercion::binary::type_union_resolution; use datafusion_expr_common::{ signature::{ArrayFunctionSignature, FIXED_SIZE_LIST_WILDCARD, TIMEZONE_WILDCARD}, @@ -873,9 +873,44 @@ fn get_valid_types( TypeSignature::Coercible(param_types) => { function_length_check(function_name, current_types.len(), param_types.len())?; + fn cast_origin( + current_type: &DataType, + encoding_preservation: EncodingPreservation, + ) -> &DataType { + match (encoding_preservation, current_type) { + ( + EncodingPreservation::Dictionary, + DataType::Dictionary(_, value_type), + ) => value_type, + _ => current_type, + } + } + + fn preserve_encoding( + current_type: &DataType, + casted_type: DataType, + encoding_preservation: EncodingPreservation, + ) -> DataType { + match (encoding_preservation, current_type, &casted_type) { + ( + EncodingPreservation::Dictionary, + DataType::Dictionary(_, _), + DataType::Dictionary(_, _), + ) => casted_type, + ( + EncodingPreservation::Dictionary, + DataType::Dictionary(key_type, _), + _, + ) => DataType::Dictionary(key_type.clone(), Box::new(casted_type)), + _ => casted_type, + } + } + let mut new_types = Vec::with_capacity(current_types.len()); for (current_type, param) in current_types.iter().zip(param_types.iter()) { let current_native_type: NativeType = current_type.into(); + let encoding_preservation = param.encoding_preservation(); + let cast_origin = cast_origin(current_type, encoding_preservation); if param .desired_type() @@ -883,9 +918,13 @@ fn get_valid_types( { let casted_type = param .desired_type() - .default_casted_type(¤t_native_type, current_type)?; + .default_casted_type(¤t_native_type, cast_origin)?; - new_types.push(casted_type); + new_types.push(preserve_encoding( + current_type, + casted_type, + encoding_preservation, + )); } else if param .allowed_source_types() .iter() @@ -894,8 +933,12 @@ fn get_valid_types( // If the condition is met which means `implicit coercion`` is provided so we can safely unwrap let default_casted_type = param.default_casted_type().unwrap(); let casted_type = - default_casted_type.default_cast_for(current_type)?; - new_types.push(casted_type); + default_casted_type.default_cast_for(cast_origin)?; + new_types.push(preserve_encoding( + current_type, + casted_type, + encoding_preservation, + )); } else { let hint = if matches!(current_native_type, NativeType::Binary) { "\n\nHint: Binary types are not automatically coerced to String. Use CAST(column AS VARCHAR) to convert Binary data to String." @@ -1214,11 +1257,11 @@ mod tests { use arrow::datatypes::IntervalUnit; use datafusion_common::{ assert_contains, - types::{logical_binary, logical_int64}, + types::{logical_binary, logical_int64, logical_string}, }; use datafusion_expr_common::{ columnar_value::ColumnarValue, - signature::{Coercion, TypeSignatureClass}, + signature::{Coercion, EncodingPreservation, TypeSignatureClass}, }; #[test] @@ -1831,6 +1874,53 @@ mod tests { Ok(()) } + #[test] + fn test_coercible_dictionary_preserves_encoding() -> Result<()> { + fn dictionary_input( + value_type: DataType, + coercion: Coercion, + ) -> Result> { + fields_with_udf( + &[Field::new( + "field", + DataType::Dictionary(Box::new(DataType::Int8), Box::new(value_type)), + true, + ) + .into()], + &MockUdf(Signature::coercible(vec![coercion], Volatility::Immutable)), + ) + .map(|v| v.into_iter().map(|f| f.data_type().clone()).collect()) + } + + let coercion = Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation(EncodingPreservation::Dictionary); + + assert_eq!( + dictionary_input(DataType::LargeUtf8, coercion.clone())?, + vec![DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::LargeUtf8), + )] + ); + assert_eq!( + dictionary_input( + DataType::BinaryView, + Coercion::new_implicit( + TypeSignatureClass::Native(logical_string()), + vec![TypeSignatureClass::Native(logical_binary())], + NativeType::String, + ) + .with_encoding_preservation(EncodingPreservation::Dictionary), + )?, + vec![DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Utf8View), + )] + ); + + Ok(()) + } + #[test] fn test_coercible_run_end_encoded() -> Result<()> { let run_end_encoded = DataType::RunEndEncoded( diff --git a/datafusion/functions/src/string/common.rs b/datafusion/functions/src/string/common.rs index 729c151f4dc3e..454302afe8f83 100644 --- a/datafusion/functions/src/string/common.rs +++ b/datafusion/functions/src/string/common.rs @@ -24,7 +24,7 @@ use crate::strings::{ StringViewArrayBuilder, append_view, }; use arrow::array::{ - Array, ArrayRef, GenericStringArray, NullBufferBuilder, OffsetSizeTrait, + Array, ArrayRef, AsArray, GenericStringArray, NullBufferBuilder, OffsetSizeTrait, StringViewArray, new_null_array, }; use arrow::buffer::{Buffer, OffsetBuffer, ScalarBuffer}; @@ -387,6 +387,9 @@ fn case_conversion( Ok(ColumnarValue::Array(Arc::new(builder.finish(nulls)?))) } + DataType::Dictionary(_, _) => Ok(ColumnarValue::Array( + case_conversion_dictionary(array, lower, name)?, + )), other => exec_err!("Unsupported data type {other:?} for function {name}"), }, ColumnarValue::Scalar(scalar) => match scalar { @@ -402,11 +405,45 @@ fn case_conversion( let result = a.as_ref().map(|x| unicode_case(x, lower)); Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(result))) } + ScalarValue::Dictionary(key_type, value) => { + let converted = case_conversion( + &[ColumnarValue::Scalar((**value).clone())], + lower, + name, + )?; + match converted { + ColumnarValue::Scalar(value) => Ok(ColumnarValue::Scalar( + ScalarValue::Dictionary(key_type.clone(), Box::new(value)), + )), + ColumnarValue::Array(_) => { + unreachable!("scalar case conversion returned an array") + } + } + } other => exec_err!("Unsupported data type {other:?} for function {name}"), }, } } +fn case_conversion_dictionary( + array: &ArrayRef, + lower: bool, + name: &str, +) -> Result { + let dictionary = array.as_any_dictionary(); + let converted = case_conversion( + &[ColumnarValue::Array(Arc::clone(dictionary.values()))], + lower, + name, + )?; + match converted { + ColumnarValue::Array(values) => Ok(dictionary.with_values(values)), + ColumnarValue::Scalar(_) => { + unreachable!("array case conversion returned a scalar") + } + } +} + fn case_conversion_array( array: &ArrayRef, lower: bool, diff --git a/datafusion/functions/src/string/lower.rs b/datafusion/functions/src/string/lower.rs index 57cbe1d8779f0..9d51942633c72 100644 --- a/datafusion/functions/src/string/lower.rs +++ b/datafusion/functions/src/string/lower.rs @@ -21,8 +21,8 @@ use crate::string::common::to_lower; use datafusion_common::Result; use datafusion_common::types::logical_string; 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; @@ -57,9 +57,10 @@ impl LowerFunc { pub fn new() -> Self { Self { signature: Signature::coercible( - vec![Coercion::new_exact(TypeSignatureClass::Native( - logical_string(), - ))], + vec![Coercion::new_exact_preserving_encoding( + TypeSignatureClass::Native(logical_string()), + EncodingPreservation::Dictionary, + )], Volatility::Immutable, ), } @@ -91,9 +92,11 @@ impl ScalarUDFImpl for LowerFunc { #[cfg(test)] mod tests { use super::*; - use arrow::array::{Array, ArrayRef, StringArray, StringViewArray}; - use arrow::datatypes::Field; - use datafusion_common::config::ConfigOptions; + use arrow::array::{ + Array, ArrayRef, DictionaryArray, Int8Array, StringArray, StringViewArray, + }; + use arrow::datatypes::{Field, Int8Type}; + use datafusion_common::{ScalarValue, config::ConfigOptions}; use std::sync::Arc; fn invoke_lower(input: ArrayRef) -> Result { @@ -118,6 +121,52 @@ mod tests { Ok(()) } + fn invoke_lower_scalar(input: ScalarValue) -> Result { + let func = LowerFunc::new(); + let data_type = input.data_type(); + let args = ScalarFunctionArgs { + number_rows: 1, + args: vec![ColumnarValue::Scalar(input)], + arg_fields: vec![Field::new("a", data_type.clone(), true).into()], + return_field: Field::new("f", data_type, true).into(), + config_options: Arc::new(ConfigOptions::default()), + }; + match func.invoke_with_args(args)? { + ColumnarValue::Scalar(result) => Ok(result), + _ => unreachable!("lower"), + } + } + + #[test] + fn lower_dictionary() -> Result<()> { + let keys = Int8Array::from(vec![Some(0), Some(1), None, Some(0)]); + let input = Arc::new(DictionaryArray::::try_new( + keys.clone(), + Arc::new(StringArray::from(vec![Some("FOO"), Some("Bar")])), + )?) as ArrayRef; + let expected = Arc::new(DictionaryArray::::try_new( + keys, + Arc::new(StringArray::from(vec![Some("foo"), Some("bar")])), + )?) as ArrayRef; + + to_lower(input, expected) + } + + #[test] + fn lower_dictionary_scalar() -> Result<()> { + let input = ScalarValue::Dictionary( + Box::new(DataType::Int8), + Box::new(ScalarValue::Utf8(Some("FOO".to_string()))), + ); + let expected = ScalarValue::Dictionary( + Box::new(DataType::Int8), + Box::new(ScalarValue::Utf8(Some("foo".to_string()))), + ); + + assert_eq!(invoke_lower_scalar(input)?, expected); + Ok(()) + } + #[test] fn lower_maybe_optimization() -> Result<()> { let input = Arc::new(StringArray::from(vec![ diff --git a/datafusion/functions/src/string/upper.rs b/datafusion/functions/src/string/upper.rs index c0ac90b1bc598..4849d7443d15c 100644 --- a/datafusion/functions/src/string/upper.rs +++ b/datafusion/functions/src/string/upper.rs @@ -20,8 +20,8 @@ use arrow::datatypes::DataType; use datafusion_common::Result; use datafusion_common::types::logical_string; 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; @@ -56,9 +56,10 @@ impl UpperFunc { pub fn new() -> Self { Self { signature: Signature::coercible( - vec![Coercion::new_exact(TypeSignatureClass::Native( - logical_string(), - ))], + vec![Coercion::new_exact_preserving_encoding( + TypeSignatureClass::Native(logical_string()), + EncodingPreservation::Dictionary, + )], Volatility::Immutable, ), } @@ -90,9 +91,11 @@ impl ScalarUDFImpl for UpperFunc { #[cfg(test)] mod tests { use super::*; - use arrow::array::{Array, ArrayRef, StringArray, StringViewArray}; - use arrow::datatypes::Field; - use datafusion_common::config::ConfigOptions; + use arrow::array::{ + Array, ArrayRef, DictionaryArray, Int8Array, StringArray, StringViewArray, + }; + use arrow::datatypes::{Field, Int8Type}; + use datafusion_common::{ScalarValue, config::ConfigOptions}; use std::sync::Arc; fn invoke_upper(input: ArrayRef) -> Result { @@ -117,6 +120,52 @@ mod tests { Ok(()) } + fn invoke_upper_scalar(input: ScalarValue) -> Result { + let func = UpperFunc::new(); + let data_type = input.data_type(); + let args = ScalarFunctionArgs { + number_rows: 1, + args: vec![ColumnarValue::Scalar(input)], + arg_fields: vec![Field::new("a", data_type.clone(), true).into()], + return_field: Field::new("f", data_type, true).into(), + config_options: Arc::new(ConfigOptions::default()), + }; + match func.invoke_with_args(args)? { + ColumnarValue::Scalar(result) => Ok(result), + _ => unreachable!("upper"), + } + } + + #[test] + fn upper_dictionary() -> Result<()> { + let keys = Int8Array::from(vec![Some(0), Some(1), None, Some(0)]); + let input = Arc::new(DictionaryArray::::try_new( + keys.clone(), + Arc::new(StringArray::from(vec![Some("foo"), Some("Bar")])), + )?) as ArrayRef; + let expected = Arc::new(DictionaryArray::::try_new( + keys, + Arc::new(StringArray::from(vec![Some("FOO"), Some("BAR")])), + )?) as ArrayRef; + + to_upper(input, expected) + } + + #[test] + fn upper_dictionary_scalar() -> Result<()> { + let input = ScalarValue::Dictionary( + Box::new(DataType::Int8), + Box::new(ScalarValue::Utf8(Some("foo".to_string()))), + ); + let expected = ScalarValue::Dictionary( + Box::new(DataType::Int8), + Box::new(ScalarValue::Utf8(Some("FOO".to_string()))), + ); + + assert_eq!(invoke_upper_scalar(input)?, expected); + Ok(()) + } + #[test] fn upper_maybe_optimization() -> Result<()> { let input = Arc::new(StringArray::from(vec![ diff --git a/datafusion/sqllogictest/test_files/functions.slt b/datafusion/sqllogictest/test_files/functions.slt index 2b393f1a26413..52aba8999f01b 100644 --- a/datafusion/sqllogictest/test_files/functions.slt +++ b/datafusion/sqllogictest/test_files/functions.slt @@ -468,7 +468,7 @@ Utf8View query T SELECT arrow_typeof(upper(arrow_cast(arrow_cast('foo', 'Dictionary(Int32, Utf8)'), 'Dictionary(Int32, Utf8View)'))) ---- -Utf8View +Dictionary(Int32, Utf8View) query T SELECT btrim(' foo ') @@ -548,7 +548,7 @@ Utf8View query T SELECT arrow_typeof(lower(arrow_cast(arrow_cast('FOObar', 'Dictionary(Int32, Utf8)'), 'Dictionary(Int32, Utf8View)'))) ---- -Utf8View +Dictionary(Int32, Utf8View) query T SELECT ltrim(' foo') From 4d32c12353f4449ada2adc74528b9a640a428d71 Mon Sep 17 00:00:00 2001 From: linfeng <33561138+lyne7-sc@users.noreply.github.com> Date: Thu, 11 Jun 2026 17:35:17 +0800 Subject: [PATCH 2/7] Update signature.rs --- datafusion/expr-common/src/signature.rs | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/datafusion/expr-common/src/signature.rs b/datafusion/expr-common/src/signature.rs index b81290bb19fc5..7c55e5514114b 100644 --- a/datafusion/expr-common/src/signature.rs +++ b/datafusion/expr-common/src/signature.rs @@ -1043,14 +1043,6 @@ fn get_data_types(native_type: &NativeType) -> Vec { /// /// * `Exact` - Only accepts arguments that exactly match the desired type /// * `Implicit` - Accepts the desired type and can coerce from specified source types -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Hash)] -pub enum EncodingPreservation { - /// Do not request preservation of a physical encoding. - None, - /// Preserve dictionary encoding and coerce only the dictionary values. - Dictionary, -} - #[derive(Debug, Clone, Eq, PartialOrd)] pub enum Coercion { /// Coercion that only accepts arguments exactly matching the desired type. @@ -1072,6 +1064,16 @@ pub enum Coercion { }, } +/// Controls whether a [`Coercion`] preserves an argument's physical encoding +/// (e.g. dictionary) instead of materializing it to the coerced value type. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Hash)] +pub enum EncodingPreservation { + /// Do not request preservation of a physical encoding. + None, + /// Preserve dictionary encoding and coerce only the dictionary values. + Dictionary, +} + impl Coercion { pub fn new_exact(desired_type: TypeSignatureClass) -> Self { Self::Exact { From 8f0f701615be232058a9fe06ac047511aad70f6e Mon Sep 17 00:00:00 2001 From: lyne7-sc <734432041@qq.com> Date: Mon, 15 Jun 2026 21:52:46 +0800 Subject: [PATCH 3/7] docs: use Coercion constructor in UDF examples --- .../library-user-guide/functions/adding-udfs.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/source/library-user-guide/functions/adding-udfs.md b/docs/source/library-user-guide/functions/adding-udfs.md index b6021c9dbb7b4..c3a40557a006d 100644 --- a/docs/source/library-user-guide/functions/adding-udfs.md +++ b/docs/source/library-user-guide/functions/adding-udfs.md @@ -397,9 +397,9 @@ impl AsyncUpper { pub fn new() -> Self { Self { signature: Signature::new( - TypeSignature::Coercible(vec![Coercion::Exact { - desired_type: TypeSignatureClass::Native(logical_string()), - }]), + TypeSignature::Coercible(vec![Coercion::new_exact( + TypeSignatureClass::Native(logical_string()), + )]), Volatility::Volatile, ), } @@ -497,9 +497,9 @@ We can now transfer the async UDF into the normal scalar using `into_scalar_udf` # pub fn new() -> Self { # Self { # signature: Signature::new( -# TypeSignature::Coercible(vec![Coercion::Exact { -# desired_type: TypeSignatureClass::Native(logical_string()), -# }]), +# TypeSignature::Coercible(vec![Coercion::new_exact( +# TypeSignatureClass::Native(logical_string()), +# )]), # Volatility::Volatile, # ), # } From b1bb38dabef002955fc3b4f2065133315bd02f64 Mon Sep 17 00:00:00 2001 From: lyne7-sc <734432041@qq.com> Date: Mon, 15 Jun 2026 22:17:48 +0800 Subject: [PATCH 4/7] use boolean flags for encoding preservation --- datafusion/expr-common/src/signature.rs | 40 ++++++++++--------- .../expr/src/type_coercion/functions.rs | 39 +++++++++--------- datafusion/functions/src/string/lower.rs | 10 +++-- datafusion/functions/src/string/upper.rs | 10 +++-- 4 files changed, 53 insertions(+), 46 deletions(-) diff --git a/datafusion/expr-common/src/signature.rs b/datafusion/expr-common/src/signature.rs index 7c55e5514114b..2e1845cac4f9d 100644 --- a/datafusion/expr-common/src/signature.rs +++ b/datafusion/expr-common/src/signature.rs @@ -1066,29 +1066,29 @@ pub enum Coercion { /// Controls whether a [`Coercion`] preserves an argument's physical encoding /// (e.g. dictionary) instead of materializing it to the coerced value type. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Hash)] -pub enum EncodingPreservation { - /// Do not request preservation of a physical encoding. - None, +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Hash)] +pub struct EncodingPreservation { + preserve_dictionary: bool, +} + +impl EncodingPreservation { /// Preserve dictionary encoding and coerce only the dictionary values. - Dictionary, + pub const fn with_dictionary(mut self) -> Self { + self.preserve_dictionary = true; + self + } + + /// Returns whether dictionary encoding should be preserved. + pub const fn preserve_dictionary(self) -> bool { + self.preserve_dictionary + } } impl Coercion { pub fn new_exact(desired_type: TypeSignatureClass) -> Self { Self::Exact { desired_type, - encoding_preservation: EncodingPreservation::None, - } - } - - pub fn new_exact_preserving_encoding( - desired_type: TypeSignatureClass, - encoding_preservation: EncodingPreservation, - ) -> Self { - Self::Exact { - desired_type, - encoding_preservation, + encoding_preservation: EncodingPreservation::default(), } } @@ -1107,7 +1107,7 @@ impl Coercion { allowed_source_types, default_casted_type, }, - encoding_preservation: EncodingPreservation::None, + encoding_preservation: EncodingPreservation::default(), } } @@ -2240,10 +2240,14 @@ mod tests { #[test] fn test_coercion_encoding_preservation_affects_equality() { + assert!(!EncodingPreservation::default().preserve_dictionary()); + let preserve_dictionary = EncodingPreservation::default().with_dictionary(); + assert!(preserve_dictionary.preserve_dictionary()); + let default = Coercion::new_exact(TypeSignatureClass::Native(logical_string())); let preserving = default .clone() - .with_encoding_preservation(EncodingPreservation::Dictionary); + .with_encoding_preservation(preserve_dictionary); assert_ne!(default, preserving); } diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs index 640de14e40bfd..f4c4152d8d393 100644 --- a/datafusion/expr/src/type_coercion/functions.rs +++ b/datafusion/expr/src/type_coercion/functions.rs @@ -877,12 +877,12 @@ fn get_valid_types( current_type: &DataType, encoding_preservation: EncodingPreservation, ) -> &DataType { - match (encoding_preservation, current_type) { - ( - EncodingPreservation::Dictionary, - DataType::Dictionary(_, value_type), - ) => value_type, - _ => current_type, + if encoding_preservation.preserve_dictionary() + && let DataType::Dictionary(_, value_type) = current_type + { + value_type + } else { + current_type } } @@ -891,18 +891,13 @@ fn get_valid_types( casted_type: DataType, encoding_preservation: EncodingPreservation, ) -> DataType { - match (encoding_preservation, current_type, &casted_type) { - ( - EncodingPreservation::Dictionary, - DataType::Dictionary(_, _), - DataType::Dictionary(_, _), - ) => casted_type, - ( - EncodingPreservation::Dictionary, - DataType::Dictionary(key_type, _), - _, - ) => DataType::Dictionary(key_type.clone(), Box::new(casted_type)), - _ => casted_type, + if encoding_preservation.preserve_dictionary() + && let DataType::Dictionary(key_type, _) = current_type + && !matches!(casted_type, DataType::Dictionary(_, _)) + { + DataType::Dictionary(key_type.clone(), Box::new(casted_type)) + } else { + casted_type } } @@ -1893,7 +1888,9 @@ mod tests { } let coercion = Coercion::new_exact(TypeSignatureClass::Native(logical_string())) - .with_encoding_preservation(EncodingPreservation::Dictionary); + .with_encoding_preservation( + EncodingPreservation::default().with_dictionary(), + ); assert_eq!( dictionary_input(DataType::LargeUtf8, coercion.clone())?, @@ -1910,7 +1907,9 @@ mod tests { vec![TypeSignatureClass::Native(logical_binary())], NativeType::String, ) - .with_encoding_preservation(EncodingPreservation::Dictionary), + .with_encoding_preservation( + EncodingPreservation::default().with_dictionary(), + ), )?, vec![DataType::Dictionary( Box::new(DataType::Int8), diff --git a/datafusion/functions/src/string/lower.rs b/datafusion/functions/src/string/lower.rs index 9d51942633c72..6e314534e688b 100644 --- a/datafusion/functions/src/string/lower.rs +++ b/datafusion/functions/src/string/lower.rs @@ -57,10 +57,12 @@ impl LowerFunc { pub fn new() -> Self { Self { signature: Signature::coercible( - vec![Coercion::new_exact_preserving_encoding( - TypeSignatureClass::Native(logical_string()), - EncodingPreservation::Dictionary, - )], + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation( + EncodingPreservation::default().with_dictionary(), + ), + ], Volatility::Immutable, ), } diff --git a/datafusion/functions/src/string/upper.rs b/datafusion/functions/src/string/upper.rs index 4849d7443d15c..c579f3a550093 100644 --- a/datafusion/functions/src/string/upper.rs +++ b/datafusion/functions/src/string/upper.rs @@ -56,10 +56,12 @@ impl UpperFunc { pub fn new() -> Self { Self { signature: Signature::coercible( - vec![Coercion::new_exact_preserving_encoding( - TypeSignatureClass::Native(logical_string()), - EncodingPreservation::Dictionary, - )], + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation( + EncodingPreservation::default().with_dictionary(), + ), + ], Volatility::Immutable, ), } From 789014d3763b9592cc5c7e7a20027d065ee7b1b5 Mon Sep 17 00:00:00 2001 From: linfeng Date: Thu, 18 Jun 2026 22:07:16 +0800 Subject: [PATCH 5/7] address dictionary preservation review feedback --- datafusion/expr-common/src/signature.rs | 9 +- .../expr/src/type_coercion/functions.rs | 38 +++- datafusion/functions/src/string/common.rs | 166 +++++++++--------- datafusion/functions/src/string/lower.rs | 58 +----- datafusion/functions/src/string/upper.rs | 58 +----- .../sqllogictest/test_files/functions.slt | 34 ++++ 6 files changed, 165 insertions(+), 198 deletions(-) diff --git a/datafusion/expr-common/src/signature.rs b/datafusion/expr-common/src/signature.rs index 2e1845cac4f9d..f0010f0a05014 100644 --- a/datafusion/expr-common/src/signature.rs +++ b/datafusion/expr-common/src/signature.rs @@ -1072,6 +1072,13 @@ pub struct EncodingPreservation { } impl EncodingPreservation { + /// Preserve dictionary encoding and coerce only the dictionary values. + pub const fn dictionary() -> Self { + Self { + preserve_dictionary: true, + } + } + /// Preserve dictionary encoding and coerce only the dictionary values. pub const fn with_dictionary(mut self) -> Self { self.preserve_dictionary = true; @@ -2241,7 +2248,7 @@ mod tests { #[test] fn test_coercion_encoding_preservation_affects_equality() { assert!(!EncodingPreservation::default().preserve_dictionary()); - let preserve_dictionary = EncodingPreservation::default().with_dictionary(); + let preserve_dictionary = EncodingPreservation::dictionary(); assert!(preserve_dictionary.preserve_dictionary()); let default = Coercion::new_exact(TypeSignatureClass::Native(logical_string())); diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs index f4c4152d8d393..12cd1f2f6e17d 100644 --- a/datafusion/expr/src/type_coercion/functions.rs +++ b/datafusion/expr/src/type_coercion/functions.rs @@ -1888,9 +1888,7 @@ mod tests { } let coercion = Coercion::new_exact(TypeSignatureClass::Native(logical_string())) - .with_encoding_preservation( - EncodingPreservation::default().with_dictionary(), - ); + .with_encoding_preservation(EncodingPreservation::dictionary()); assert_eq!( dictionary_input(DataType::LargeUtf8, coercion.clone())?, @@ -1907,15 +1905,43 @@ mod tests { vec![TypeSignatureClass::Native(logical_binary())], NativeType::String, ) - .with_encoding_preservation( - EncodingPreservation::default().with_dictionary(), - ), + .with_encoding_preservation(EncodingPreservation::dictionary()), )?, vec![DataType::Dictionary( Box::new(DataType::Int8), Box::new(DataType::Utf8View), )] ); + assert_eq!( + dictionary_input( + DataType::Int32, + Coercion::new_implicit( + TypeSignatureClass::Native(logical_int64()), + vec![TypeSignatureClass::Integer], + NativeType::Int64, + ) + .with_encoding_preservation(EncodingPreservation::dictionary()), + )?, + vec![DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Int64), + )] + ); + assert_eq!( + dictionary_input( + DataType::Int32, + Coercion::new_implicit( + TypeSignatureClass::Integer, + vec![], + NativeType::Int64, + ) + .with_encoding_preservation(EncodingPreservation::dictionary()), + )?, + vec![DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Int32), + )] + ); Ok(()) } diff --git a/datafusion/functions/src/string/common.rs b/datafusion/functions/src/string/common.rs index 454302afe8f83..b51b92e9df1ed 100644 --- a/datafusion/functions/src/string/common.rs +++ b/datafusion/functions/src/string/common.rs @@ -348,81 +348,90 @@ fn case_conversion( name: &str, ) -> Result { match &args[0] { - ColumnarValue::Array(array) => match array.data_type() { - DataType::Utf8 => Ok(ColumnarValue::Array(case_conversion_array::( - array, lower, - )?)), - DataType::LargeUtf8 => Ok(ColumnarValue::Array( - case_conversion_array::(array, lower)?, - )), - DataType::Utf8View => { - let string_array = as_string_view_array(array)?; - if string_array.is_ascii() { - return Ok(ColumnarValue::Array(Arc::new( - case_conversion_utf8view_ascii(string_array, lower), - ))); - } - let item_len = string_array.len(); - // Null-preserving: reuse the input null buffer as the output null buffer. - let nulls = string_array.nulls().cloned(); - let mut builder = StringViewArrayBuilder::with_capacity(item_len); - - if let Some(ref n) = nulls { - for i in 0..item_len { - if n.is_null(i) { - builder.try_append_placeholder()?; - } else { - // SAFETY: `n.is_null(i)` was false in the branch above. - let s = unsafe { string_array.value_unchecked(i) }; - builder.try_append_value(&unicode_case(s, lower))?; - } - } - } else { - for i in 0..item_len { - // SAFETY: no null buffer means every index is valid. - let s = unsafe { string_array.value_unchecked(i) }; - builder.try_append_value(&unicode_case(s, lower))?; - } - } + ColumnarValue::Array(array) => Ok(ColumnarValue::Array( + case_conversion_columnar_array(array, lower, name)?, + )), + ColumnarValue::Scalar(scalar) => Ok(ColumnarValue::Scalar( + case_conversion_scalar(scalar, lower, name)?, + )), + } +} - Ok(ColumnarValue::Array(Arc::new(builder.finish(nulls)?))) - } - DataType::Dictionary(_, _) => Ok(ColumnarValue::Array( - case_conversion_dictionary(array, lower, name)?, - )), - other => exec_err!("Unsupported data type {other:?} for function {name}"), - }, - ColumnarValue::Scalar(scalar) => match scalar { - ScalarValue::Utf8(a) => { - let result = a.as_ref().map(|x| unicode_case(x, lower)); - Ok(ColumnarValue::Scalar(ScalarValue::Utf8(result))) - } - ScalarValue::LargeUtf8(a) => { - let result = a.as_ref().map(|x| unicode_case(x, lower)); - Ok(ColumnarValue::Scalar(ScalarValue::LargeUtf8(result))) - } - ScalarValue::Utf8View(a) => { - let result = a.as_ref().map(|x| unicode_case(x, lower)); - Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(result))) - } - ScalarValue::Dictionary(key_type, value) => { - let converted = case_conversion( - &[ColumnarValue::Scalar((**value).clone())], - lower, - name, - )?; - match converted { - ColumnarValue::Scalar(value) => Ok(ColumnarValue::Scalar( - ScalarValue::Dictionary(key_type.clone(), Box::new(value)), - )), - ColumnarValue::Array(_) => { - unreachable!("scalar case conversion returned an array") - } - } +fn case_conversion_scalar( + scalar: &ScalarValue, + lower: bool, + name: &str, +) -> Result { + match scalar { + ScalarValue::Utf8(a) => { + let result = a.as_ref().map(|x| unicode_case(x, lower)); + Ok(ScalarValue::Utf8(result)) + } + ScalarValue::LargeUtf8(a) => { + let result = a.as_ref().map(|x| unicode_case(x, lower)); + Ok(ScalarValue::LargeUtf8(result)) + } + ScalarValue::Utf8View(a) => { + let result = a.as_ref().map(|x| unicode_case(x, lower)); + Ok(ScalarValue::Utf8View(result)) + } + ScalarValue::Dictionary(key_type, value) => { + let converted = case_conversion_scalar(value.as_ref(), lower, name)?; + Ok(ScalarValue::Dictionary( + key_type.clone(), + Box::new(converted), + )) + } + other => exec_err!("Unsupported data type {other:?} for function {name}"), + } +} + +fn case_conversion_columnar_array( + array: &ArrayRef, + lower: bool, + name: &str, +) -> Result { + match array.data_type() { + DataType::Utf8 => case_conversion_array::(array, lower), + DataType::LargeUtf8 => case_conversion_array::(array, lower), + DataType::Utf8View => case_conversion_utf8view(array, lower), + DataType::Dictionary(_, _) => case_conversion_dictionary(array, lower, name), + other => exec_err!("Unsupported data type {other:?} for function {name}"), + } +} + +fn case_conversion_utf8view(array: &ArrayRef, lower: bool) -> Result { + let string_array = as_string_view_array(array)?; + if string_array.is_ascii() { + return Ok(Arc::new(case_conversion_utf8view_ascii( + string_array, + lower, + ))); + } + let item_len = string_array.len(); + // Null-preserving: reuse the input null buffer as the output null buffer. + let nulls = string_array.nulls().cloned(); + let mut builder = StringViewArrayBuilder::with_capacity(item_len); + + if let Some(ref n) = nulls { + for i in 0..item_len { + if n.is_null(i) { + builder.try_append_placeholder()?; + } else { + // SAFETY: `n.is_null(i)` was false in the branch above. + let s = unsafe { string_array.value_unchecked(i) }; + builder.try_append_value(&unicode_case(s, lower))?; } - other => exec_err!("Unsupported data type {other:?} for function {name}"), - }, + } + } else { + for i in 0..item_len { + // SAFETY: no null buffer means every index is valid. + let s = unsafe { string_array.value_unchecked(i) }; + builder.try_append_value(&unicode_case(s, lower))?; + } } + + Ok(Arc::new(builder.finish(nulls)?)) } fn case_conversion_dictionary( @@ -431,17 +440,8 @@ fn case_conversion_dictionary( name: &str, ) -> Result { let dictionary = array.as_any_dictionary(); - let converted = case_conversion( - &[ColumnarValue::Array(Arc::clone(dictionary.values()))], - lower, - name, - )?; - match converted { - ColumnarValue::Array(values) => Ok(dictionary.with_values(values)), - ColumnarValue::Scalar(_) => { - unreachable!("array case conversion returned a scalar") - } - } + let converted = case_conversion_columnar_array(dictionary.values(), lower, name)?; + Ok(dictionary.with_values(converted)) } fn case_conversion_array( diff --git a/datafusion/functions/src/string/lower.rs b/datafusion/functions/src/string/lower.rs index 6e314534e688b..88f2c800e9e0c 100644 --- a/datafusion/functions/src/string/lower.rs +++ b/datafusion/functions/src/string/lower.rs @@ -59,9 +59,7 @@ impl LowerFunc { signature: Signature::coercible( vec![ Coercion::new_exact(TypeSignatureClass::Native(logical_string())) - .with_encoding_preservation( - EncodingPreservation::default().with_dictionary(), - ), + .with_encoding_preservation(EncodingPreservation::dictionary()), ], Volatility::Immutable, ), @@ -94,11 +92,9 @@ impl ScalarUDFImpl for LowerFunc { #[cfg(test)] mod tests { use super::*; - use arrow::array::{ - Array, ArrayRef, DictionaryArray, Int8Array, StringArray, StringViewArray, - }; - use arrow::datatypes::{Field, Int8Type}; - use datafusion_common::{ScalarValue, config::ConfigOptions}; + use arrow::array::{Array, ArrayRef, StringArray, StringViewArray}; + use arrow::datatypes::Field; + use datafusion_common::config::ConfigOptions; use std::sync::Arc; fn invoke_lower(input: ArrayRef) -> Result { @@ -123,52 +119,6 @@ mod tests { Ok(()) } - fn invoke_lower_scalar(input: ScalarValue) -> Result { - let func = LowerFunc::new(); - let data_type = input.data_type(); - let args = ScalarFunctionArgs { - number_rows: 1, - args: vec![ColumnarValue::Scalar(input)], - arg_fields: vec![Field::new("a", data_type.clone(), true).into()], - return_field: Field::new("f", data_type, true).into(), - config_options: Arc::new(ConfigOptions::default()), - }; - match func.invoke_with_args(args)? { - ColumnarValue::Scalar(result) => Ok(result), - _ => unreachable!("lower"), - } - } - - #[test] - fn lower_dictionary() -> Result<()> { - let keys = Int8Array::from(vec![Some(0), Some(1), None, Some(0)]); - let input = Arc::new(DictionaryArray::::try_new( - keys.clone(), - Arc::new(StringArray::from(vec![Some("FOO"), Some("Bar")])), - )?) as ArrayRef; - let expected = Arc::new(DictionaryArray::::try_new( - keys, - Arc::new(StringArray::from(vec![Some("foo"), Some("bar")])), - )?) as ArrayRef; - - to_lower(input, expected) - } - - #[test] - fn lower_dictionary_scalar() -> Result<()> { - let input = ScalarValue::Dictionary( - Box::new(DataType::Int8), - Box::new(ScalarValue::Utf8(Some("FOO".to_string()))), - ); - let expected = ScalarValue::Dictionary( - Box::new(DataType::Int8), - Box::new(ScalarValue::Utf8(Some("foo".to_string()))), - ); - - assert_eq!(invoke_lower_scalar(input)?, expected); - Ok(()) - } - #[test] fn lower_maybe_optimization() -> Result<()> { let input = Arc::new(StringArray::from(vec![ diff --git a/datafusion/functions/src/string/upper.rs b/datafusion/functions/src/string/upper.rs index c579f3a550093..789ab2c046203 100644 --- a/datafusion/functions/src/string/upper.rs +++ b/datafusion/functions/src/string/upper.rs @@ -58,9 +58,7 @@ impl UpperFunc { signature: Signature::coercible( vec![ Coercion::new_exact(TypeSignatureClass::Native(logical_string())) - .with_encoding_preservation( - EncodingPreservation::default().with_dictionary(), - ), + .with_encoding_preservation(EncodingPreservation::dictionary()), ], Volatility::Immutable, ), @@ -93,11 +91,9 @@ impl ScalarUDFImpl for UpperFunc { #[cfg(test)] mod tests { use super::*; - use arrow::array::{ - Array, ArrayRef, DictionaryArray, Int8Array, StringArray, StringViewArray, - }; - use arrow::datatypes::{Field, Int8Type}; - use datafusion_common::{ScalarValue, config::ConfigOptions}; + use arrow::array::{Array, ArrayRef, StringArray, StringViewArray}; + use arrow::datatypes::Field; + use datafusion_common::config::ConfigOptions; use std::sync::Arc; fn invoke_upper(input: ArrayRef) -> Result { @@ -122,52 +118,6 @@ mod tests { Ok(()) } - fn invoke_upper_scalar(input: ScalarValue) -> Result { - let func = UpperFunc::new(); - let data_type = input.data_type(); - let args = ScalarFunctionArgs { - number_rows: 1, - args: vec![ColumnarValue::Scalar(input)], - arg_fields: vec![Field::new("a", data_type.clone(), true).into()], - return_field: Field::new("f", data_type, true).into(), - config_options: Arc::new(ConfigOptions::default()), - }; - match func.invoke_with_args(args)? { - ColumnarValue::Scalar(result) => Ok(result), - _ => unreachable!("upper"), - } - } - - #[test] - fn upper_dictionary() -> Result<()> { - let keys = Int8Array::from(vec![Some(0), Some(1), None, Some(0)]); - let input = Arc::new(DictionaryArray::::try_new( - keys.clone(), - Arc::new(StringArray::from(vec![Some("foo"), Some("Bar")])), - )?) as ArrayRef; - let expected = Arc::new(DictionaryArray::::try_new( - keys, - Arc::new(StringArray::from(vec![Some("FOO"), Some("BAR")])), - )?) as ArrayRef; - - to_upper(input, expected) - } - - #[test] - fn upper_dictionary_scalar() -> Result<()> { - let input = ScalarValue::Dictionary( - Box::new(DataType::Int8), - Box::new(ScalarValue::Utf8(Some("foo".to_string()))), - ); - let expected = ScalarValue::Dictionary( - Box::new(DataType::Int8), - Box::new(ScalarValue::Utf8(Some("FOO".to_string()))), - ); - - assert_eq!(invoke_upper_scalar(input)?, expected); - Ok(()) - } - #[test] fn upper_maybe_optimization() -> Result<()> { let input = Arc::new(StringArray::from(vec![ diff --git a/datafusion/sqllogictest/test_files/functions.slt b/datafusion/sqllogictest/test_files/functions.slt index 52aba8999f01b..98edfa189d3e3 100644 --- a/datafusion/sqllogictest/test_files/functions.slt +++ b/datafusion/sqllogictest/test_files/functions.slt @@ -470,6 +470,23 @@ SELECT arrow_typeof(upper(arrow_cast(arrow_cast('foo', 'Dictionary(Int32, Utf8)' ---- Dictionary(Int32, Utf8View) +statement ok +CREATE TABLE upper_dictionary_test AS +SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS c1 FROM (VALUES +('foo'), +(NULL), +('Bar')); + +query TT +SELECT upper(c1), arrow_typeof(upper(c1)) FROM upper_dictionary_test +---- +FOO Dictionary(Int32, Utf8) +NULL Dictionary(Int32, Utf8) +BAR Dictionary(Int32, Utf8) + +statement ok +DROP TABLE upper_dictionary_test + query T SELECT btrim(' foo ') ---- @@ -550,6 +567,23 @@ SELECT arrow_typeof(lower(arrow_cast(arrow_cast('FOObar', 'Dictionary(Int32, Utf ---- Dictionary(Int32, Utf8View) +statement ok +CREATE TABLE lower_dictionary_test AS +SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS c1 FROM (VALUES +('FOO'), +(NULL), +('Bar')); + +query TT +SELECT lower(c1), arrow_typeof(lower(c1)) FROM lower_dictionary_test +---- +foo Dictionary(Int32, Utf8) +NULL Dictionary(Int32, Utf8) +bar Dictionary(Int32, Utf8) + +statement ok +DROP TABLE lower_dictionary_test + query T SELECT ltrim(' foo') ---- From adbdd13f816f68a5f016e17e9dce9f59fddf3720 Mon Sep 17 00:00:00 2001 From: linfeng <33561138+lyne7-sc@users.noreply.github.com> Date: Thu, 18 Jun 2026 23:01:17 +0800 Subject: [PATCH 6/7] add tests --- .../expr/src/type_coercion/functions.rs | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs index 12cd1f2f6e17d..c2dc56ae1008a 100644 --- a/datafusion/expr/src/type_coercion/functions.rs +++ b/datafusion/expr/src/type_coercion/functions.rs @@ -1912,6 +1912,19 @@ mod tests { Box::new(DataType::Utf8View), )] ); + // Contrast: without encoding_preservation, Native strips dictionary entirely + assert_eq!( + dictionary_input( + DataType::Int32, + Coercion::new_implicit( + TypeSignatureClass::Native(logical_int64()), + vec![TypeSignatureClass::Integer], + NativeType::Int64, + ), + )?, + vec![DataType::Int64] + ); + // With encoding_preservation, dictionary wrapper is preserved, value coerced assert_eq!( dictionary_input( DataType::Int32, @@ -1927,6 +1940,22 @@ mod tests { Box::new(DataType::Int64), )] ); + // Contrast: without encoding_preservation, non-Native already passes through + assert_eq!( + dictionary_input( + DataType::Int32, + Coercion::new_implicit( + TypeSignatureClass::Integer, + vec![], + NativeType::Int64, + ), + )?, + vec![DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Int32), + )] + ); + // With encoding_preservation, same result — no difference for non-Native assert_eq!( dictionary_input( DataType::Int32, From 364d9d083af7e6b206e4f93bde07d8aee820defb Mon Sep 17 00:00:00 2001 From: linfeng <33561138+lyne7-sc@users.noreply.github.com> Date: Tue, 7 Jul 2026 14:31:37 +0800 Subject: [PATCH 7/7] added upgrading documentation --- .../library-user-guide/upgrading/55.0.0.md | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index 307de722bd13e..de45ce902c5a7 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -82,6 +82,27 @@ will now appear as `Decimal128(NULL,10,2)`. Query result values already used human-readable decimal formatting and are unchanged. +### `Coercion` supports dictionary encoding preservation + +`datafusion_expr_common::signature::Coercion` now supports optional dictionary +encoding preservation. When enabled for `TypeSignatureClass::Native(...)` +coercions, DataFusion coerces dictionary inputs to +`Dictionary(original_key_type, coerced_value_type)` instead of materializing them +to the coerced value type. + +User-defined functions can opt in by setting dictionary encoding preservation on +the relevant coercion: + +```rust +Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation(EncodingPreservation::dictionary()) +``` + +This changes the coerced argument type passed to the function. If a function +derives its return type from that coerced argument type, code that checks exact +result types may need to update its expectations or add an explicit cast to +materialize the result. + ### `GroupsAccumulator::merge_batch` no longer takes `opt_filter` The `opt_filter` argument has been removed from