You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by gi...@apache.org on 2026/10/18 01:32:09 UTC
(arrow-rs) branch gh-readonly-queue/main/pr-11502-6db2d4fd99c766ecd65a90419cea717ad5fbb619 created (now 5effeb5123)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-11502-6db2d4fd99c766ecd65a90419cea717ad5fbb619
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
at 5effeb5123 speed up dict cast key width (#11502)
This branch includes the following new commits:
new 5effeb5123 speed up dict cast key width (#11502)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(arrow-rs) 01/01: speed up dict cast key width (#11502)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11502-6db2d4fd99c766ecd65a90419cea717ad5fbb619
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit 5effeb51234f713c4cc81070b038f9dfc58e6cfd
Author: RIchard Baah <13...@users.noreply.github.com>
AuthorDate: Sun Oct 18 01:31:45 2026 +0000
speed up dict cast key width (#11502)
# Which issue does this PR close?
<!--
It is not necessary to raise an issue, but it is encouraged if it is a
non-trivial
enhancement/bug.
If this PR relates to an existing issue, but does not fully resolve it,
please change
the wording below from "Closes #NNN" to "Relates to #NNN" or "Part of
#NNN" to avoid
false closing an issue.
-->
- works towards https://github.com/apache/datafusion/issues/24111
- for https://github.com/apache/datafusion/issues/24117 we need to
reduce the cost of selecting a key type that is too small for a given
cardinality. this means making it as cheap as possible to bump the key
type of a dictionary array.
# Rationale for this change
`Dictionary(K1, V) → Dictionary(K2, V)` casts (changing only the key
width) currently take the general path in
`dictionary_to_dictionary_cast`, which does three things that are
unnecessary when the value type is unchanged:
1. Calls `cast_with_options` on the values array — dispatch overhead
even for identity casts, real work for near-identity casts.
2. Rebuilds the dictionary's `ArrayData` graph via
`into_builder().child_data(...).build_unchecked()`.
3. Wraps the result in `DictionaryArray::<K2>::from(data)`, which runs
an `O(N_rows)` bounds-check pass over every key.
**If the source dict was valid and the key cast didn't overflow, the
result is valid by construction**
<!--
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.
Please try to explain the problem you are trying to solve in terms of
the user-visible
behavior, rather than the implementation.
-->
# What changes are included in this PR?
A fast path in `dictionary_to_dictionary_cast`: when the source value
type equals the target value type (and the source isn't a nested dict),
delegate to `cast_keys_only`, which casts the keys array and reuses the
values `Arc` via `DictionaryArray::new_unchecked`. Narrowing overflow
remains a hard error regardless of `cast_options.safe`, same contract as
the general path.
<!--
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.
-->
# Are these changes tested?
<!--
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, if so identify these tests)?
If this PR claims a performance improvement, please include evidence
such as
benchmark results.
DO NOT simply paste a list of commands ran to verify the code; this is
what our
CI is for and it pollutes the PR. ONLY explain if testing was included
with this
PR, or why it wasn't.
-->
# Are there any significant user-facing changes?
no
<!--
Significant user-facing changes include:
- Breaking public API change
- Significant change to existing functionality (e.g. function returns
something different)
These need to be explicitly called out to ensure we have proper
documentation if
required and to ensure we follow proper semantic versioning.
DO NOT restate the changes of this PR; if a bug was fixed and already
explained
above, we don't need this to be restated as a user-facing change.
-->
---
arrow-cast/src/cast/dictionary.rs | 59 +++++++++++++++++++++++++++++++++++++
arrow-cast/tests/cast/dictionary.rs | 22 ++++++++++++--
2 files changed, 79 insertions(+), 2 deletions(-)
diff --git a/arrow-cast/src/cast/dictionary.rs b/arrow-cast/src/cast/dictionary.rs
index 25e247e16f..181b4c4f30 100644
--- a/arrow-cast/src/cast/dictionary.rs
+++ b/arrow-cast/src/cast/dictionary.rs
@@ -132,6 +132,13 @@ fn dictionary_to_dictionary_cast<K: ArrowDictionaryKeyType>(
) -> Result<ArrayRef, ArrowError> {
use DataType::*;
+ // Value type unchanged (and not nested dict): cast keys only, reuse values Arc.
+ if array.values().data_type() == to_value_type
+ && !matches!(array.values().data_type(), Dictionary(_, _))
+ {
+ return cast_keys_only(array, to_index_type, cast_options);
+ }
+
// Fast path for a nested dictionary source (`Dictionary<K, Dictionary<K2, V>>`).
// Both layers index into the same inner values, so the two index layers can
// be composed into one rather than materializing the values: `take` gathers
@@ -199,6 +206,58 @@ fn dictionary_to_dictionary_cast<K: ArrowDictionaryKeyType>(
Ok(new_array)
}
+fn finish_keys_only<K2: ArrowDictionaryKeyType>(
+ cast_keys: ArrayRef,
+ values: ArrayRef,
+) -> Result<ArrayRef, ArrowError> {
+ let new_keys = cast_keys
+ .as_any()
+ .downcast_ref::<PrimitiveArray<K2>>()
+ .ok_or_else(|| {
+ ArrowError::ComputeError(
+ "Internal Error: cast keys did not downcast to target key type".to_string(),
+ )
+ })?
+ .clone();
+ // SAFETY: caller rejects overflowing keys, so every non-null key stays in bounds of `values`.
+ Ok(Arc::new(unsafe {
+ DictionaryArray::<K2>::new_unchecked(new_keys, values)
+ }))
+}
+
+fn cast_keys_only<K: ArrowDictionaryKeyType>(
+ array: &DictionaryArray<K>,
+ to_index_type: &DataType,
+ cast_options: &CastOptions,
+) -> Result<ArrayRef, ArrowError> {
+ let source_keys = array.keys();
+ let cast_keys = cast_with_options(&source_keys, to_index_type, cast_options)?;
+
+ if cast_keys.null_count() > source_keys.null_count() {
+ return Err(ArrowError::ComputeError(format!(
+ "Could not convert {} dictionary indexes from {:?} to {:?}",
+ cast_keys.null_count() - source_keys.null_count(),
+ source_keys.data_type(),
+ to_index_type
+ )));
+ }
+
+ let values = array.values().clone();
+ match to_index_type {
+ DataType::Int8 => finish_keys_only::<Int8Type>(cast_keys, values),
+ DataType::Int16 => finish_keys_only::<Int16Type>(cast_keys, values),
+ DataType::Int32 => finish_keys_only::<Int32Type>(cast_keys, values),
+ DataType::Int64 => finish_keys_only::<Int64Type>(cast_keys, values),
+ DataType::UInt8 => finish_keys_only::<UInt8Type>(cast_keys, values),
+ DataType::UInt16 => finish_keys_only::<UInt16Type>(cast_keys, values),
+ DataType::UInt32 => finish_keys_only::<UInt32Type>(cast_keys, values),
+ DataType::UInt64 => finish_keys_only::<UInt64Type>(cast_keys, values),
+ _ => Err(ArrowError::CastError(format!(
+ "Unsupported type {to_index_type} for dictionary index"
+ ))),
+ }
+}
+
/// Cast `Dict<K, Binary>` or `Dict<K, LargeBinary>` to `Utf8View`, validating UTF-8 for each
/// dictionary value.
///
diff --git a/arrow-cast/tests/cast/dictionary.rs b/arrow-cast/tests/cast/dictionary.rs
index ef5ee193b4..13ee6e80ea 100644
--- a/arrow-cast/tests/cast/dictionary.rs
+++ b/arrow-cast/tests/cast/dictionary.rs
@@ -25,8 +25,8 @@ use arrow_array::cast::AsArray;
use arrow_array::types::{Int8Type, Int32Type, Int64Type, UInt16Type, UInt32Type};
use arrow_array::{
Array, ArrayRef, BinaryArray, BinaryViewArray, Date32Array, DictionaryArray,
- FixedSizeBinaryArray, Int32Array, LargeStringArray, StringArray, StringViewArray, StructArray,
- TimestampSecondArray,
+ FixedSizeBinaryArray, Int8Array, Int32Array, LargeStringArray, StringArray, StringViewArray,
+ StructArray, TimestampSecondArray,
};
use arrow_buffer::NullBuffer;
use arrow_cast::display::{ArrayFormatter, FormatOptions};
@@ -673,6 +673,24 @@ fn test_cast_dict_to_dict_bad_index_value_utf8() {
);
}
+#[test]
+fn test_cast_dict_to_dict_key_width_only_reuses_values() {
+ let values = Arc::new(StringArray::from(vec!["foo", "bar"])) as ArrayRef;
+ let dict = DictionaryArray::<Int8Type>::new(
+ Int8Array::from(vec![Some(0), None, Some(1), Some(0)]),
+ values.clone(),
+ );
+
+ let result = cast(&dict, &Dictionary(Box::new(Int32), Box::new(Utf8))).unwrap();
+ let result = result.as_dictionary::<Int32Type>();
+
+ assert_eq!(
+ result.keys(),
+ &Int32Array::from(vec![Some(0), None, Some(1), Some(0)])
+ );
+ assert!(Arc::ptr_eq(&values, result.values()));
+}
+
#[test]
fn test_cast_nested_dictionary_to_dictionary_reuses_values() {
let inner = DictionaryArray::<Int32Type>::new(