Skip to content

perf: preserve dictionary encoding for lower/upper to avoid materializing low-cardinality columns #22905

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lyne7-sc wants to merge 7 commits into
base: main
Choose a base branch
from
Open

perf: preserve dictionary encoding for lower/upper to avoid materializing low-cardinality columns #22905

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
1896e06
perf: preserve dictionary encoding for case conversion
lyne7-sc Jun 10, 2026
4d32c12
Update signature.rs
lyne7-sc Jun 11, 2026
8f0f701
docs: use Coercion constructor in UDF examples
lyne7-sc Jun 15, 2026
b1bb38d
use boolean flags for encoding preservation
lyne7-sc Jun 15, 2026
789014d
address dictionary preservation review feedback
lyne7-sc Jun 18, 2026
adbdd13
add tests
lyne7-sc Jun 18, 2026
364d9d0
added upgrading documentation
lyne7-sc Jul 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 83 additions & 2 deletions datafusion/expr-common/src/signature.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
}
}
Comment on lines +1067 to +1092

@Jefffrey Jefffrey Jun 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Something to consider is if an enum is best suited for this 🤔

For example, if we want to also include run encoded arrays as well (since they are similar to dictionaries), would this mean two new variants, one just for run arrays and one for dictionary + run arrays?

  • I don't know if arrow is planning to include any more types of encodings like this at the moment

So I was also thinking perhaps a bitflag approach (or just a struct with boolean flags) could be another approach:

struct Encoding {
    preserve_dictionary: bool,
    preserve_run: bool,
}

But I don't know how common a use case would be to implement preservation only for dictionaries and not run arrays (or vice versa) so maybe thats overengineering 🤔

@lyne7-sc lyne7-sc Jun 15, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Considering RunEndEncoded and other potential encoding types, I agree that using a struct with boolean flags is a better approach. It keeps the complexity low while making it easier to add future preservation options without introducing enum variants for every possible combination.


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.
Expand All @@ -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,
}
}

Expand All @@ -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,
}
}
Expand All @@ -1128,13 +1193,15 @@ 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()
}
}

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);
}
}

Expand Down Expand Up @@ -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;
Expand Down
4 changes: 2 additions & 2 deletions datafusion/expr/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down
158 changes: 151 additions & 7 deletions datafusion/expr/src/type_coercion/functions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -873,19 +873,53 @@ 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()
.matches_native_type(&current_native_type)
{
let casted_type = param
.desired_type()
.default_casted_type(&current_native_type, current_type)?;
.default_casted_type(&current_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()
Expand All @@ -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."
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -1831,6 +1869,112 @@ mod tests {
Ok(())
}

#[test]
fn test_coercible_dictionary_preserves_encoding() -> Result<()> {

@Jefffrey Jefffrey Jun 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we also have a test for a TypeSignatureClass that isn't native? curious to see how it'll work, since #19458 highlighted that current behaviour for dictionaries is already different for TypeSignatureClass::Native and non-native (e.g. TypeSignatureClass::Integer)

@lyne7-sc lyne7-sc Jun 18, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added tests for both cases:

The behavior now looks like this:

TypeSignatureClass no preservation dictionary preservation
Native(Int64) Int64 Dictionary(Int8, Int64)
Non-Native(Integer) Dictionary(Int8, Int32) Dictionary(Int8, Int32)

@Jefffrey Jefffrey Jul 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for checking this; I'm a little worried because of the discrepancy here now; is it a big task to try align them? Though we might need to check the UDFs of existing functions that use a non-native typesignatureclass in case they already knew this assumption and had dictionary paths 🤔

@lyne7-sc lyne7-sc Jul 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks like this discrepancy already exists today. Typed non-Native TypeSignatureClass values can pass Dictionary through, while Native coercions materialize to the target type. But this pr makes that contrast more visible by adding an explicit preservation API for the Native path.

If we want EncodingPreservation to be the general mechanism for deciding whether Dictionary is preserved, then I agree we should also think about typed non-Native TypeSignatureClass values. I took a quick pass over the affected built-ins. The good news is that I did not find an obvious core DataFusion function in the affected typed non-Native TypeSignatureClass path that intentionally depends on receiving Dictionary(_, T)🙂. Most seem to dispatch on T directly, and some would likely reject Dictionary(_, T) today after type coercion has accepted it. From that angle, aligning the default behavior may actually fix cases where dictionary inputs currently pass signature matching but fail later in the function implementation.

I did find a couple of Spark compatibility functions, such as spark hex and bitmap_count, that have dictionary paths and would likely need explicit opt-in if we align the default behavior.

From my side, the set of built-ins that may need changes looks bounded, so aligning the behavior seems manageable.

@Jefffrey Jefffrey Jul 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for checking this; I think we can separate that work into a followup PR 👍

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(
Expand Down
Loading
Loading