This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 3089ace491 perf: preserve dictionary encoding for lower/upper to avoid
materializing low-cardinality columns (#22905)
3089ace491 is described below
commit 3089ace491701b1b1d0cb047df66d4821dfdf15e
Author: linfeng <[email protected]>
AuthorDate: Wed Jul 8 23:29:15 2026 +0800
perf: preserve dictionary encoding for lower/upper to avoid materializing
low-cardinality columns (#22905)
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->
- Related to #19458
- Related to #20935
## Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
When a `Dictionary(K, Utf8)` column is passed to a string scalar
function, the type-coercion layer currently materializes it to flat
Utf8/Utf8View before the function runs, so the operation is applied to
every row instead of just the unique dictionary values, and the
dictionary encoding is lost on the output. See #19458 for the underlying
coercion behavior and #20935 for the string-function-specific impact.
This is wasteful for low-cardinality columns and inflates Arrow
IPC/Flight message sizes downstream.
## What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
- Add `EncodingPreservation { None, Dictionary }` and opt-in
constructors on `Coercion` (`new_exact_preserving_encoding`,
`with_encoding_preservation`).
- In `get_valid_types`, when a `Coercible` arg requests `Dictionary`
preservation, run coercion against the dictionary's value type and
re-wrap the result as `Dictionary(K, V')`, so the function receives a
`DictionaryArray`.
- Make `lower/upper` opt in and handle dictionary inputs (array +
scalar) by converting only the dictionary values and re-wrapping with
the original keys.
## Are these changes tested?
Yes
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->
## Are there any user-facing changes?
Yes. upper/lower now return Dictionary(...) for dictionary inputs
instead of the previously materialized Utf8View. New public API
(EncodingPreservation, new Coercion constructors) is added — please add
the api change label.
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->
---
datafusion/expr-common/src/signature.rs | 85 ++++++++++-
datafusion/expr/src/lib.rs | 4 +-
datafusion/expr/src/type_coercion/functions.rs | 158 ++++++++++++++++++++-
datafusion/functions/src/string/common.rs | 147 ++++++++++++-------
datafusion/functions/src/string/lower.rs | 11 +-
datafusion/functions/src/string/upper.rs | 11 +-
datafusion/sqllogictest/test_files/functions.slt | 38 ++++-
.../library-user-guide/functions/adding-udfs.md | 12 +-
docs/source/library-user-guide/upgrading/55.0.0.md | 21 +++
9 files changed, 403 insertions(+), 84 deletions(-)
diff --git a/datafusion/expr-common/src/signature.rs
b/datafusion/expr-common/src/signature.rs
index 35a679cc44..f0010f0a05 100644
--- a/datafusion/expr-common/src/signature.rs
+++ b/datafusion/expr-common/src/signature.rs
@@ -1049,6 +1049,8 @@ pub enum Coercion {
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 +1059,44 @@ 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,
},
}
+/// 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, Default, PartialEq, Eq, PartialOrd, Hash)]
+pub struct EncodingPreservation {
+ preserve_dictionary: bool,
+}
+
+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;
+ 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 }
+ Self::Exact {
+ desired_type,
+ encoding_preservation: EncodingPreservation::default(),
+ }
}
/// Create a new coercion with implicit coercion rules.
@@ -1080,6 +1114,37 @@ impl Coercion {
allowed_source_types,
default_casted_type,
},
+ encoding_preservation: EncodingPreservation::default(),
+ }
+ }
+
+ 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 +1168,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 +1193,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 +1201,7 @@ impl Hash for Coercion {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.desired_type().hash(state);
self.implicit_coercion().hash(state);
+ self.encoding_preservation().hash(state);
}
}
@@ -2178,6 +2245,20 @@ mod tests {
assert_snapshot!(implicit_with_multiple_sources, @"Int64");
}
+ #[test]
+ fn test_coercion_encoding_preservation_affects_equality() {
+ assert!(!EncodingPreservation::default().preserve_dictionary());
+ let preserve_dictionary = EncodingPreservation::dictionary();
+ assert!(preserve_dictionary.preserve_dictionary());
+
+ let default =
Coercion::new_exact(TypeSignatureClass::Native(logical_string()));
+ let preserving = default
+ .clone()
+ .with_encoding_preservation(preserve_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 b52a784df9..43cb3fdc20 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 33746a2c46..c2dc56ae10 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,39 @@ 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 {
+ if encoding_preservation.preserve_dictionary()
+ && let DataType::Dictionary(_, value_type) = current_type
+ {
+ value_type
+ } else {
+ current_type
+ }
+ }
+
+ fn preserve_encoding(
+ current_type: &DataType,
+ casted_type: DataType,
+ encoding_preservation: EncodingPreservation,
+ ) -> DataType {
+ 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
+ }
+ }
+
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 +913,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 +928,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 +1252,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 +1869,112 @@ mod tests {
Ok(())
}
+ #[test]
+ fn test_coercible_dictionary_preserves_encoding() -> Result<()> {
+ fn dictionary_input(
+ value_type: DataType,
+ coercion: Coercion,
+ ) -> Result<Vec<DataType>> {
+ 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),
+ )]
+ );
+ // 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,
+ 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),
+ )]
+ );
+ // 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,
+ 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(())
+ }
+
#[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 729c151f4d..b51b92e9df 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};
@@ -348,63 +348,100 @@ fn case_conversion(
name: &str,
) -> Result<ColumnarValue> {
match &args[0] {
- ColumnarValue::Array(array) => match array.data_type() {
- DataType::Utf8 =>
Ok(ColumnarValue::Array(case_conversion_array::<i32>(
- array, lower,
- )?)),
- DataType::LargeUtf8 => Ok(ColumnarValue::Array(
- case_conversion_array::<i64>(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)?)))
- }
- 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)))
+fn case_conversion_scalar(
+ scalar: &ScalarValue,
+ lower: bool,
+ name: &str,
+) -> Result<ScalarValue> {
+ 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<ArrayRef> {
+ match array.data_type() {
+ DataType::Utf8 => case_conversion_array::<i32>(array, lower),
+ DataType::LargeUtf8 => case_conversion_array::<i64>(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<ArrayRef>
{
+ 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(
+ array: &ArrayRef,
+ lower: bool,
+ name: &str,
+) -> Result<ArrayRef> {
+ let dictionary = array.as_any_dictionary();
+ let converted = case_conversion_columnar_array(dictionary.values(), lower,
name)?;
+ Ok(dictionary.with_values(converted))
}
fn case_conversion_array<O: OffsetSizeTrait>(
diff --git a/datafusion/functions/src/string/lower.rs
b/datafusion/functions/src/string/lower.rs
index 57cbe1d877..88f2c800e9 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(TypeSignatureClass::Native(logical_string()))
+
.with_encoding_preservation(EncodingPreservation::dictionary()),
+ ],
Volatility::Immutable,
),
}
diff --git a/datafusion/functions/src/string/upper.rs
b/datafusion/functions/src/string/upper.rs
index c0ac90b1bc..789ab2c046 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(TypeSignatureClass::Native(logical_string()))
+
.with_encoding_preservation(EncodingPreservation::dictionary()),
+ ],
Volatility::Immutable,
),
}
diff --git a/datafusion/sqllogictest/test_files/functions.slt
b/datafusion/sqllogictest/test_files/functions.slt
index 2b393f1a26..98edfa189d 100644
--- a/datafusion/sqllogictest/test_files/functions.slt
+++ b/datafusion/sqllogictest/test_files/functions.slt
@@ -468,7 +468,24 @@ Utf8View
query T
SELECT arrow_typeof(upper(arrow_cast(arrow_cast('foo', 'Dictionary(Int32,
Utf8)'), 'Dictionary(Int32, Utf8View)')))
----
-Utf8View
+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 ')
@@ -548,7 +565,24 @@ Utf8View
query T
SELECT arrow_typeof(lower(arrow_cast(arrow_cast('FOObar', 'Dictionary(Int32,
Utf8)'), 'Dictionary(Int32, Utf8View)')))
----
-Utf8View
+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')
diff --git a/docs/source/library-user-guide/functions/adding-udfs.md
b/docs/source/library-user-guide/functions/adding-udfs.md
index b6021c9dbb..c3a40557a0 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,
# ),
# }
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 f39b34a640..26a26b69b5 100644
--- a/docs/source/library-user-guide/upgrading/55.0.0.md
+++ b/docs/source/library-user-guide/upgrading/55.0.0.md
@@ -124,6 +124,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
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
