This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new fbe75a3bd9 arrow-cast: Add optimized path for unnesting a dict (#10248)
fbe75a3bd9 is described below
commit fbe75a3bd9df6d66658e2954e14163e5a93a4187
Author: Frederic Branczyk <[email protected]>
AuthorDate: Tue Jun 30 20:33:49 2026 +0200
arrow-cast: Add optimized path for unnesting a dict (#10248)
# Which issue does this PR close?
- Closes https://github.com/apache/arrow-rs/issues/10247
# What changes are included in this PR?
Benchmark and optimization.
# Are these changes tested?
Benchmark:
`cargo bench -p arrow --bench cast_kernels --features test_utils --
"cast nested dict to dict"`
Before (unpack inner values):
```
cast nested dict to dict
time: [12.148 µs 12.217 µs 12.283 µs]
```
After (compose indices, reuse inner values):
```
cast nested dict to dict
time: [2.7624 µs 2.7679 µs 2.7754 µs]
change: [−77.294% −77.149% −77.001%] (p = 0.00 <
0.05)
Performance has improved.
```
~4.4x faster
# Are there any user-facing changes?
No, just an optimization.
@alamb @Jefffrey
---
arrow-cast/src/cast/dictionary.rs | 18 ++++++++++++++++++
arrow-cast/src/cast/mod.rs | 30 ++++++++++++++++++++++++++++++
arrow/benches/cast_kernels.rs | 28 ++++++++++++++++++++++++++++
3 files changed, 76 insertions(+)
diff --git a/arrow-cast/src/cast/dictionary.rs
b/arrow-cast/src/cast/dictionary.rs
index 83aa691482..db3611f551 100644
--- a/arrow-cast/src/cast/dictionary.rs
+++ b/arrow-cast/src/cast/dictionary.rs
@@ -59,6 +59,24 @@ fn dictionary_to_dictionary_cast<K: ArrowDictionaryKeyType>(
) -> Result<ArrayRef, ArrowError> {
use DataType::*;
+ // 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
+ // the inner keys through the outer keys and reuses the inner values buffer
+ // untouched, so no value data is rewritten. The flattened single-level
+ // dictionary is then cast to the requested index/value types.
+ if matches!(array.values().data_type(), Dictionary(_, _)) {
+ let flattened = take(array.values().as_ref(), array.keys(), None)?;
+ return cast_with_options(
+ &flattened,
+ &Dictionary(
+ Box::new(to_index_type.clone()),
+ Box::new(to_value_type.clone()),
+ ),
+ cast_options,
+ );
+ }
+
let keys_array: ArrayRef =
Arc::new(PrimitiveArray::<K>::from(array.keys().to_data()));
let values_array = array.values();
let cast_keys = cast_with_options(&keys_array, to_index_type,
cast_options)?;
diff --git a/arrow-cast/src/cast/mod.rs b/arrow-cast/src/cast/mod.rs
index d903bd0f04..20a3068e95 100644
--- a/arrow-cast/src/cast/mod.rs
+++ b/arrow-cast/src/cast/mod.rs
@@ -8932,6 +8932,36 @@ mod tests {
);
}
+ #[test]
+ fn test_cast_nested_dictionary_to_dictionary_reuses_values() {
+ let inner = DictionaryArray::<Int32Type>::new(
+ Int32Array::from(vec![Some(0), None, Some(1)]),
+ Arc::new(StringArray::from(vec!["x", "y"])),
+ );
+ let nested = DictionaryArray::<Int32Type>::new(
+ Int32Array::from(vec![Some(0), Some(1), Some(2), None, Some(0)]),
+ Arc::new(inner),
+ );
+
+ let result = cast(&nested, &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), None, Some(0)])
+ );
+ assert_eq!(
+ result.values().as_string::<i32>(),
+ &StringArray::from(vec!["x", "y"])
+ );
+ let logical: Vec<Option<&str>> = result
+ .downcast_dict::<StringArray>()
+ .unwrap()
+ .into_iter()
+ .collect();
+ assert_eq!(logical, vec![Some("x"), None, Some("y"), None, Some("x")]);
+ }
+
#[test]
fn test_cast_primitive_dict() {
// FROM a dictionary with of INT32 values
diff --git a/arrow/benches/cast_kernels.rs b/arrow/benches/cast_kernels.rs
index 82e9696f04..7e98168924 100644
--- a/arrow/benches/cast_kernels.rs
+++ b/arrow/benches/cast_kernels.rs
@@ -185,6 +185,25 @@ fn build_dict_array(size: usize) -> ArrayRef {
Arc::new(DictionaryArray::new(keys, Arc::new(values)))
}
+// Dictionary<UInt32, Dictionary<UInt32, Utf8>>: an inner dictionary of 4096
+// entries over 64 moderately long distinct values, wrapped by `size` outer
keys
+// that index into it. Casting to Dictionary<UInt32, Utf8> flattens the two
index
+// layers; the old path unpacks (copies) the inner values, the fast path reuses
+// them and only remaps indices.
+fn build_nested_dict_array(size: usize) -> ArrayRef {
+ let distinct: Vec<String> = (0..64)
+ .map(|i| format!("a moderately long symbol name number {i:04}"))
+ .collect();
+ let inner_values = StringArray::from_iter(distinct.iter().map(|s|
Some(s.as_str())));
+ let inner_len = 4096u32;
+ let distinct_len = inner_values.len() as u32;
+ let inner_keys = UInt32Array::from_iter((0..inner_len).map(|v| v %
distinct_len));
+ let inner = DictionaryArray::new(inner_keys, Arc::new(inner_values));
+
+ let outer_keys = UInt32Array::from_iter((0..size as u32).map(|v| v %
inner_len));
+ Arc::new(DictionaryArray::new(outer_keys, Arc::new(inner)))
+}
+
// cast array from specified primitive array type to desired data type
fn cast_array(array: &ArrayRef, to_type: DataType) {
hint::black_box(cast(hint::black_box(array),
hint::black_box(&to_type)).unwrap());
@@ -212,6 +231,7 @@ fn add_benchmark(c: &mut Criterion) {
let wide_string_array = cast(&string_array, &DataType::LargeUtf8).unwrap();
let dict_array = build_dict_array(10_000);
+ let nested_dict_array = build_nested_dict_array(10_000);
let string_view_array = cast(&dict_array, &DataType::Utf8View).unwrap();
let binary_view_array = cast(&string_view_array,
&DataType::BinaryView).unwrap();
@@ -330,6 +350,14 @@ fn add_benchmark(c: &mut Criterion) {
c.bench_function("cast dict to string view", |b| {
b.iter(|| cast_array(&dict_array, DataType::Utf8View))
});
+ c.bench_function("cast nested dict to dict", |b| {
+ b.iter(|| {
+ cast_array(
+ &nested_dict_array,
+ DataType::Dictionary(Box::new(DataType::UInt32),
Box::new(DataType::Utf8)),
+ )
+ })
+ });
c.bench_function("cast string view to dict", |b| {
b.iter(|| {
cast_array(
