Skip to content

fix(arrow-data): don't panic on dictionary key overflow in interleave/concat #10675

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
okhsunrog wants to merge 1 commit into
base: main
Choose a base branch
from
+149 −20
Open

fix(arrow-data): don't panic on dictionary key overflow in interleave/concat #10675

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
68 changes: 50 additions & 18 deletions arrow-data/src/transform/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 `Err(ArrowError::DictionaryKeyOverflowError)`, 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.
///
Expand All @@ -418,12 +432,30 @@ 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.
/// of `arrays`. Or when a [Capacities] variant is not yet supported. Or 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 `Err(ArrowError::DictionaryKeyOverflowError)` 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) {
Expand Down Expand Up @@ -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![],
Expand All @@ -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;
Expand All @@ -554,19 +586,19 @@ 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| {
let child_arrays = arrays
.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
Expand All @@ -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) => {
Expand All @@ -596,19 +628,19 @@ 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| {
let child_arrays = arrays
.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
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -716,15 +748,15 @@ impl<'a> MutableArrayData<'a> {
buffer2,
child_data,
};
Self {
Ok(Self {
arrays,
data,
dictionary,
variadic_data_buffers,
extend_values,
extend_null_bits,
extend_nulls,
}
})
}

/// Extends the in progress array with a region of the input arrays, returning an error on
Expand Down
47 changes: 46 additions & 1 deletion arrow-select/src/concat.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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())?
Expand Down Expand Up @@ -1732,6 +1732,51 @@ mod tests {
);
}

#[test]
fn concat_string_view_dictionary_overflow_returns_err() {
// Two independently-built `Dictionary<UInt8, Utf8View>` arrays, each within
// the u8 key range on its own, but whose *combined* distinct values overflow
// it (analogous to per-partition dictionary-encoded columns being merged).
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));

// Must not panic: the key type genuinely cannot address 400 distinct values,
// so this should surface as a normal, catchable error.
let err = concat(&[&dict_a, &dict_b]).unwrap_err();
assert!(matches!(err, ArrowError::DictionaryKeyOverflowError));
}

#[test]
fn concat_nested_dictionary_overflow_returns_err() {
// Same overflow as `concat_string_view_dictionary_overflow_returns_err`, but
// with the dictionary nested inside a `FixedSizeList`, exercising the
// recursive child construction in `MutableArrayData::try_with_capacities`
// rather than the top-level dictionary handling.
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() {
Expand Down
54 changes: 53 additions & 1 deletion arrow-select/src/interleave.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -2033,6 +2033,58 @@ mod tests {
);
}

#[test]
fn test_interleave_string_view_dictionary_overflow_returns_err() {
// Two independently-built `Dictionary<UInt8, Utf8View>` arrays, each within
// the u8 key range on its own, but whose *combined* distinct values overflow
// it. This mirrors what happens when a UInt16-keyed dictionary column is
// scanned in multiple partitions and merged (e.g. via a sort-preserving
// merge), each partition building its own dictionary independently.
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();

// Must not panic: the key type genuinely cannot address 400 distinct values,
// so this should surface as a normal, catchable error.
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 overflow as `test_interleave_string_view_dictionary_overflow_returns_err`,
// but with the dictionary nested inside a `FixedSizeList`, exercising the
// recursive child construction in `MutableArrayData::try_with_capacities`
// (reached via `interleave_fallback`) rather than the top-level dictionary
// handling.
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() {
Expand Down