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 365c76066c fix(arrow-data): don't panic on dictionary key overflow in
interleave/concat (#10675)
365c76066c is described below
commit 365c76066c1ca647090eaf61b4b4df536cca82e1
Author: Danila Gornushko <[email protected]>
AuthorDate: Thu Aug 20 04:31:40 2026 +0300
fix(arrow-data): don't panic on dictionary key overflow in
interleave/concat (#10675)
# Which issue does this PR close?
- Closes #10674.
# Rationale for this change
`interleave()` and `concat()` document a `Result` return type, but for
`Dictionary<K, Utf8View>`/`Dictionary<K, BinaryView>` arrays (top-level
or nested inside a
`List`/`FixedSizeList`/`Struct`/`RunEndEncoded`/`Union`), a genuine
dictionary key overflow currently panics instead of returning `Err`.
This happens because the fallback path builds a `MutableArrayData`
directly and `.expect()`s the dictionary-concat result, and
`MutableArrayData::new`/`with_capacities` have no fallible variant. See
#10674 for the full analysis and a minimal repro.
# What changes are included in this PR?
- Add `MutableArrayData::try_new` / `try_with_capacities`, fallible
variants of `new` / `with_capacities` that return
`Err(ArrowError::DictionaryKeyOverflowError)` instead of panicking on
dictionary key overflow. `new()`/`with_capacities()` keep their existing
panicking behavior (now implemented as
`try_with_capacities(..).expect(..)`) for the many existing callers that
rely on infallibility.
- Recursive child construction inside `try_with_capacities` (for
`List`/`LargeList`/`Map`/`ListView`/`LargeListView`/`FixedSizeList`/`Struct`/`RunEndEncoded`/`Union`)
also uses the fallible variants and propagates errors with `?`, so
dictionaries nested inside container types are covered, not just
top-level dictionary arrays.
- Switch `arrow-select`'s `interleave_fallback` and `concat_fallback`
(the paths reached for dictionary arrays whose values can't/shouldn't be
merged) to the fallible constructors.
# Are these changes tested?
Yes:
- `concat_string_view_dictionary_overflow_returns_err` /
`test_interleave_string_view_dictionary_overflow_returns_err`: top-level
`Dictionary<UInt8, Utf8View>` overflow returns `Err` instead of
panicking.
- `concat_nested_dictionary_overflow_returns_err` /
`test_interleave_nested_dictionary_overflow_returns_err`: same overflow
nested inside a `FixedSizeList`, exercising the recursive child
construction.
Full `arrow-data`/`arrow-select` test suites pass (410 tests), `cargo
fmt --check` and `cargo clippy --all-targets -- -D warnings` are clean
for both crates.
# Are there any user-facing changes?
No breaking changes. `MutableArrayData::new`/`with_capacities` keep
their documented panicking behavior and signatures. Two new public
fallible methods are added (`try_new`, `try_with_capacities`).
`interleave()`/`concat()` keep their existing `Result` signature -- the
only visible change is that a specific previously-panicking input
(genuine dictionary key overflow on `Utf8View`/`BinaryView`
dictionaries) now returns `Err(ArrowError::DictionaryKeyOverflowError)`
instead.
---------
Co-authored-by: Jeffrey Vo <[email protected]>
---
arrow-data/src/transform/mod.rs | 70 ++++++++++++++++++++++++++++++-----------
arrow-select/src/concat.rs | 41 +++++++++++++++++++++++-
arrow-select/src/interleave.rs | 45 +++++++++++++++++++++++++-
3 files changed, 135 insertions(+), 21 deletions(-)
diff --git a/arrow-data/src/transform/mod.rs b/arrow-data/src/transform/mod.rs
index 7145463ac4..6571e0b658 100644
--- a/arrow-data/src/transform/mod.rs
+++ b/arrow-data/src/transform/mod.rs
@@ -410,6 +410,20 @@ impl<'a> MutableArrayData<'a> {
Self::with_capacities(arrays, use_nulls, Capacities::Array(capacity))
}
+ /// Fallible variant of [MutableArrayData::new].
+ ///
+ /// Unlike [MutableArrayData::new], this does not panic when merging
dictionary
+ /// arrays whose combined values would overflow the dictionary key type.
Instead,
+ /// it returns an error, letting callers (e.g. [`interleave`](crate) /
`concat`)
+ /// surface it as a normal error.
+ pub fn try_new(
+ arrays: Vec<&'a ArrayData>,
+ use_nulls: bool,
+ capacity: usize,
+ ) -> Result<Self, ArrowError> {
+ Self::try_with_capacities(arrays, use_nulls,
Capacities::Array(capacity))
+ }
+
/// Similar to [MutableArrayData::new], but lets users define the
/// preallocated capacities of the array with more granularity.
///
@@ -417,13 +431,31 @@ impl<'a> MutableArrayData<'a> {
///
/// # Panics
///
- /// This function panics if the given `capacities` don't match the data
type
- /// of `arrays`. Or when a [Capacities] variant is not yet supported.
+ /// * if the given `capacities` don't match the data type of `arrays`
+ /// * if a [Capacities] variant is not yet supported
+ /// * when merging dictionary arrays whose combined values overflow the
+ /// dictionary key type — see [MutableArrayData::try_with_capacities]
for a
+ /// fallible variant
pub fn with_capacities(
arrays: Vec<&'a ArrayData>,
use_nulls: bool,
capacities: Capacities,
) -> Self {
+ Self::try_with_capacities(arrays, use_nulls, capacities)
+ .expect("MutableArrayData::new is infallible")
+ }
+
+ /// Fallible variant of [MutableArrayData::with_capacities].
+ ///
+ /// Returns an error instead of panicking when merging dictionary arrays
whose
+ /// combined values would overflow the dictionary key type. Still panics
for
+ /// other unsupported combinations (inconsistent input types, unsupported
+ /// `Capacities` variants) as documented on
[MutableArrayData::with_capacities].
+ pub fn try_with_capacities(
+ arrays: Vec<&'a ArrayData>,
+ use_nulls: bool,
+ capacities: Capacities,
+ ) -> Result<Self, ArrowError> {
let data_type = arrays[0].data_type();
for a in arrays.iter().skip(1) {
@@ -522,9 +554,9 @@ impl<'a> MutableArrayData<'a> {
Capacities::Array(array_capacity)
};
- vec![MutableArrayData::with_capacities(
+ vec![MutableArrayData::try_with_capacities(
children, use_nulls, capacities,
- )]
+ )?]
}
// the dictionary type just appends keys and clones the values.
DataType::Dictionary(_, _) => vec![],
@@ -538,13 +570,13 @@ impl<'a> MutableArrayData<'a> {
.iter()
.map(|array| &array.child_data()[i])
.collect::<Vec<_>>();
- MutableArrayData::with_capacities(
+ MutableArrayData::try_with_capacities(
child_arrays,
use_nulls,
child_cap.clone(),
)
})
- .collect::<Vec<_>>()
+ .collect::<Result<Vec<_>, _>>()?
}
Capacities::Struct(capacity, None) => {
array_capacity = capacity;
@@ -554,9 +586,9 @@ impl<'a> MutableArrayData<'a> {
.iter()
.map(|array| &array.child_data()[i])
.collect::<Vec<_>>();
- MutableArrayData::new(child_arrays, use_nulls,
capacity)
+ MutableArrayData::try_new(child_arrays, use_nulls,
capacity)
})
- .collect::<Vec<_>>()
+ .collect::<Result<Vec<_>, _>>()?
}
_ => (0..fields.len())
.map(|i| {
@@ -564,9 +596,9 @@ impl<'a> MutableArrayData<'a> {
.iter()
.map(|array| &array.child_data()[i])
.collect::<Vec<_>>();
- MutableArrayData::new(child_arrays, use_nulls,
array_capacity)
+ MutableArrayData::try_new(child_arrays, use_nulls,
array_capacity)
})
- .collect::<Vec<_>>(),
+ .collect::<Result<Vec<_>, _>>()?,
},
DataType::RunEndEncoded(_, _) => {
let run_ends_child = arrays
@@ -578,8 +610,8 @@ impl<'a> MutableArrayData<'a> {
.map(|array| &array.child_data()[1])
.collect::<Vec<_>>();
vec![
- MutableArrayData::new(run_ends_child, false,
array_capacity),
- MutableArrayData::new(value_child, use_nulls,
array_capacity),
+ MutableArrayData::try_new(run_ends_child, false,
array_capacity)?,
+ MutableArrayData::try_new(value_child, use_nulls,
array_capacity)?,
]
}
DataType::FixedSizeList(_, size) => {
@@ -596,9 +628,9 @@ impl<'a> MutableArrayData<'a> {
} else {
Capacities::Array(array_capacity * *size as usize)
};
- vec![MutableArrayData::with_capacities(
+ vec![MutableArrayData::try_with_capacities(
children, use_nulls, capacities,
- )]
+ )?]
}
DataType::Union(fields, _) => (0..fields.len())
.map(|i| {
@@ -606,9 +638,9 @@ impl<'a> MutableArrayData<'a> {
.iter()
.map(|array| &array.child_data()[i])
.collect::<Vec<_>>();
- MutableArrayData::new(child_arrays, use_nulls,
array_capacity)
+ MutableArrayData::try_new(child_arrays, use_nulls,
array_capacity)
})
- .collect::<Vec<_>>(),
+ .collect::<Result<Vec<_>, _>>()?,
};
// Get the dictionary if any, and if it is a concatenation of multiple
@@ -688,7 +720,7 @@ impl<'a> MutableArrayData<'a> {
})
.collect();
- extend_values.expect("MutableArrayData::new is infallible")
+ extend_values?
}
DataType::BinaryView | DataType::Utf8View => {
let mut next_offset = 0u32;
@@ -716,7 +748,7 @@ impl<'a> MutableArrayData<'a> {
buffer2,
child_data,
};
- Self {
+ Ok(Self {
arrays,
data,
dictionary,
@@ -724,7 +756,7 @@ impl<'a> MutableArrayData<'a> {
extend_values,
extend_null_bits,
extend_nulls,
- }
+ })
}
/// Extends the in progress array with a region of the input arrays,
returning an error on
diff --git a/arrow-select/src/concat.rs b/arrow-select/src/concat.rs
index a8d9117415..e1384edbc5 100644
--- a/arrow-select/src/concat.rs
+++ b/arrow-select/src/concat.rs
@@ -579,7 +579,7 @@ pub fn concat(arrays: &[&dyn Array]) -> Result<ArrayRef,
ArrowError> {
fn concat_fallback(arrays: &[&dyn Array], capacity: Capacities) ->
Result<ArrayRef, ArrowError> {
let array_data: Vec<_> = arrays.iter().map(|a|
a.to_data()).collect::<Vec<_>>();
let array_data = array_data.iter().collect();
- let mut mutable = MutableArrayData::with_capacities(array_data, false,
capacity);
+ let mut mutable = MutableArrayData::try_with_capacities(array_data, false,
capacity)?;
for (i, a) in arrays.iter().enumerate() {
mutable.try_extend(i, 0, a.len())?
@@ -1732,6 +1732,45 @@ mod tests {
);
}
+ #[test]
+ fn concat_string_view_dictionary_overflow_returns_err() {
+ // concatenating dictionaries which results in overflowing the key
type should
+ // surface an error not a panic
+ let values_a: StringViewArray = (0..200).map(|i|
Some(format!("a{i}"))).collect();
+ let keys_a = UInt8Array::from_iter_values(0..200);
+ let dict_a = DictionaryArray::<UInt8Type>::new(keys_a,
Arc::new(values_a));
+
+ let values_b: StringViewArray = (0..200).map(|i|
Some(format!("b{i}"))).collect();
+ let keys_b = UInt8Array::from_iter_values(0..200);
+ let dict_b = DictionaryArray::<UInt8Type>::new(keys_b,
Arc::new(values_b));
+
+ let err = concat(&[&dict_a, &dict_b]).unwrap_err();
+ assert!(matches!(err, ArrowError::DictionaryKeyOverflowError));
+ }
+
+ #[test]
+ fn concat_nested_dictionary_overflow_returns_err() {
+ // same as above, but with the dictionary nested inside a FixedSizeList
+ let field = Arc::new(arrow_schema::Field::new(
+ "item",
+ DataType::Dictionary(Box::new(DataType::UInt8),
Box::new(DataType::Utf8View)),
+ false,
+ ));
+
+ let values_a: StringViewArray = (0..200).map(|i|
Some(format!("a{i}"))).collect();
+ let keys_a = UInt8Array::from_iter_values(0..200);
+ let dict_a = DictionaryArray::<UInt8Type>::new(keys_a,
Arc::new(values_a));
+ let list_a = FixedSizeListArray::new(field.clone(), 1,
Arc::new(dict_a), None);
+
+ let values_b: StringViewArray = (0..200).map(|i|
Some(format!("b{i}"))).collect();
+ let keys_b = UInt8Array::from_iter_values(0..200);
+ let dict_b = DictionaryArray::<UInt8Type>::new(keys_b,
Arc::new(values_b));
+ let list_b = FixedSizeListArray::new(field, 1, Arc::new(dict_b), None);
+
+ let err = concat(&[&list_a, &list_b]).unwrap_err();
+ assert!(matches!(err, ArrowError::DictionaryKeyOverflowError));
+ }
+
#[test]
#[cfg_attr(miri, ignore)] // Takes too long
fn concat_many_dictionary_list_arrays() {
diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs
index f56864ac65..6d5033b1d4 100644
--- a/arrow-select/src/interleave.rs
+++ b/arrow-select/src/interleave.rs
@@ -768,7 +768,7 @@ fn interleave_fallback(
) -> Result<ArrayRef, ArrowError> {
let arrays: Vec<_> = values.iter().map(|x| x.to_data()).collect();
let arrays: Vec<_> = arrays.iter().collect();
- let mut array_data = MutableArrayData::new(arrays, false, indices.len());
+ let mut array_data = MutableArrayData::try_new(arrays, false,
indices.len())?;
let mut cur_array = indices[0].0;
let mut start_row_idx = indices[0].1;
@@ -2033,6 +2033,49 @@ mod tests {
);
}
+ #[test]
+ fn test_interleave_string_view_dictionary_overflow_returns_err() {
+ // interleaving dictionaries which results in overflowing the key type
should
+ // surface an error not a panic
+ let values_a: StringViewArray = (0..200).map(|i|
Some(format!("a{i}"))).collect();
+ let keys_a = UInt8Array::from_iter_values(0..200);
+ let dict_a = DictionaryArray::<UInt8Type>::new(keys_a,
Arc::new(values_a));
+
+ let values_b: StringViewArray = (0..200).map(|i|
Some(format!("b{i}"))).collect();
+ let keys_b = UInt8Array::from_iter_values(0..200);
+ let dict_b = DictionaryArray::<UInt8Type>::new(keys_b,
Arc::new(values_b));
+
+ let indices: Vec<_> = (0..200).flat_map(|i| [(0, i), (1,
i)]).collect();
+
+ let err = interleave(&[&dict_a, &dict_b], &indices).unwrap_err();
+ assert!(matches!(err, ArrowError::DictionaryKeyOverflowError));
+ }
+
+ #[test]
+ fn test_interleave_nested_dictionary_overflow_returns_err() {
+ // same as above, but with the dictionary nested inside a FixedSizeList
+ let field = Arc::new(arrow_schema::Field::new(
+ "item",
+ DataType::Dictionary(Box::new(DataType::UInt8),
Box::new(DataType::Utf8View)),
+ false,
+ ));
+
+ let values_a: StringViewArray = (0..200).map(|i|
Some(format!("a{i}"))).collect();
+ let keys_a = UInt8Array::from_iter_values(0..200);
+ let dict_a = DictionaryArray::<UInt8Type>::new(keys_a,
Arc::new(values_a));
+ let list_a = FixedSizeListArray::new(field.clone(), 1,
Arc::new(dict_a), None);
+
+ let values_b: StringViewArray = (0..200).map(|i|
Some(format!("b{i}"))).collect();
+ let keys_b = UInt8Array::from_iter_values(0..200);
+ let dict_b = DictionaryArray::<UInt8Type>::new(keys_b,
Arc::new(values_b));
+ let list_b = FixedSizeListArray::new(field, 1, Arc::new(dict_b), None);
+
+ let indices: Vec<_> = (0..200).flat_map(|i| [(0, i), (1,
i)]).collect();
+
+ let err = interleave(&[&list_a, &list_b], &indices).unwrap_err();
+ assert!(matches!(err, ArrowError::DictionaryKeyOverflowError));
+ }
+
#[test]
#[cfg_attr(miri, ignore)] // Takes too long
fn test_interleave_bytes_offset_overflow() {
