This is an automated email from the ASF dual-hosted git repository.
Jefffrey 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 2edc0a044f fix(arrow-data): allow full dictionary key range when
concatenating (#10323)
2edc0a044f is described below
commit 2edc0a044f495606e65878dadffe8d1a64f43c01
Author: raphaelroshan <[email protected]>
AuthorDate: Thu Jul 16 01:21:21 2026 +0800
fix(arrow-data): allow full dictionary key range when concatenating (#10323)
# Which issue does this PR close?
Closes #9366.
# Rationale for this change
Concatenating dictionary arrays whose merged dictionary length is
exactly the key type's maximum value + 1 (e.g. 256 distinct values with
`u8` keys) panics with `DictionaryKeyOverflowError`. The largest key
index is `length - 1`, so 256 values use keys `0..=255`, which fit in
`u8` — the overflow check was off by one.
# What changes are included in this PR?
`build_extend_dictionary` now checks that the key type can hold `max -
1` (the largest index) rather than `max` (the length).
# Are these changes tested?
Yes — added `test_dict_overflow_9366`, which concatenates two
`Dictionary<u8, _>` arrays totaling 256 values and asserts it succeeds
(panics before this change).
# Are there any user-facing changes?
No API change; a previously-panicking concat now succeeds.
Co-authored-by: raphaelroshan <[email protected]>
---
arrow-data/src/transform/mod.rs | 5 ++++-
arrow-select/src/concat.rs | 50 +++++++++++++++++++++++++++++++++++++++++
2 files changed, 54 insertions(+), 1 deletion(-)
diff --git a/arrow-data/src/transform/mod.rs b/arrow-data/src/transform/mod.rs
index 21570d8c02..7b43c090a9 100644
--- a/arrow-data/src/transform/mod.rs
+++ b/arrow-data/src/transform/mod.rs
@@ -195,7 +195,10 @@ impl std::fmt::Debug for MutableArrayData<'_> {
fn build_extend_dictionary(array: &ArrayData, offset: usize, max: usize) ->
Option<Extend<'_>> {
macro_rules! validate_and_build {
($dt: ty) => {{
- let _: $dt = max.try_into().ok()?;
+ // `max` is the merged dictionary length; the largest key index is
+ // `max - 1`, so the key type only needs to hold `max - 1` (e.g.
256
+ // values use keys 0..=255, which fit in u8).
+ let _: $dt = max.saturating_sub(1).try_into().ok()?;
let offset: $dt = offset.try_into().ok()?;
Some(primitive::build_extend_with_offset(array, offset))
}};
diff --git a/arrow-select/src/concat.rs b/arrow-select/src/concat.rs
index 839fceae34..f18624c5ad 100644
--- a/arrow-select/src/concat.rs
+++ b/arrow-select/src/concat.rs
@@ -644,6 +644,56 @@ mod tests {
use arrow_schema::{Field, Schema};
use std::fmt::Debug;
+ #[test]
+ fn test_dict_overflow_9366() {
+ use arrow_schema::DataType;
+
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "a",
+ DataType::Dictionary(
+ Box::new(DataType::UInt8),
+ Box::new(DataType::FixedSizeBinary(8)),
+ ),
+ false,
+ )]));
+ let make = |vals: std::ops::Range<u64>| {
+ let dict = FixedSizeBinaryArray::try_from_iter(vals.map(|i|
i.to_le_bytes())).unwrap();
+ let keys = UInt8Array::from_iter_values(0..128);
+ let arr = DictionaryArray::try_new(keys, Arc::new(dict)).unwrap();
+ RecordBatch::try_new(schema.clone(), vec![Arc::new(arr)]).unwrap()
+ };
+ // 256 distinct values fit in u8 keys (0..=255): concat must succeed.
+ let out = concat_batches(&schema, &[make(0..128),
make(128..256)]).unwrap();
+ assert_eq!(out.num_rows(), 256);
+ let dict = out.column(0).as_dictionary::<UInt8Type>();
+ assert_eq!(dict.values().len(), 256);
+ }
+
+ #[test]
+ fn test_dict_overflow_i8_9366() {
+ use arrow_schema::DataType;
+
+ // Same boundary for a signed key type: i8 holds 128 keys (0..=127).
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "a",
+ DataType::Dictionary(
+ Box::new(DataType::Int8),
+ Box::new(DataType::FixedSizeBinary(8)),
+ ),
+ false,
+ )]));
+ let make = |vals: std::ops::Range<u64>| {
+ let dict = FixedSizeBinaryArray::try_from_iter(vals.map(|i|
i.to_le_bytes())).unwrap();
+ let keys = Int8Array::from_iter_values(0..64);
+ let arr = DictionaryArray::try_new(keys, Arc::new(dict)).unwrap();
+ RecordBatch::try_new(schema.clone(), vec![Arc::new(arr)]).unwrap()
+ };
+ let out = concat_batches(&schema, &[make(0..64),
make(64..128)]).unwrap();
+ assert_eq!(out.num_rows(), 128);
+ let dict = out.column(0).as_dictionary::<Int8Type>();
+ assert_eq!(dict.values().len(), 128);
+ }
+
#[test]
fn test_concat_empty_vec() {
let re = concat(&[]);
