diff --git a/arrow-data/src/transform/mod.rs b/arrow-data/src/transform/mod.rs index f57d8fcc9b94..5210167390d6 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 `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::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. /// @@ -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 { 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::>(); - MutableArrayData::with_capacities( + MutableArrayData::try_with_capacities( child_arrays, use_nulls, child_cap.clone(), ) }) - .collect::>() + .collect::, _>>()? } Capacities::Struct(capacity, None) => { array_capacity = capacity; @@ -554,9 +586,9 @@ impl<'a> MutableArrayData<'a> { .iter() .map(|array| &array.child_data()[i]) .collect::>(); - MutableArrayData::new(child_arrays, use_nulls, capacity) + MutableArrayData::try_new(child_arrays, use_nulls, capacity) }) - .collect::>() + .collect::, _>>()? } _ => (0..fields.len()) .map(|i| { @@ -564,9 +596,9 @@ impl<'a> MutableArrayData<'a> { .iter() .map(|array| &array.child_data()[i]) .collect::>(); - MutableArrayData::new(child_arrays, use_nulls, array_capacity) + MutableArrayData::try_new(child_arrays, use_nulls, array_capacity) }) - .collect::>(), + .collect::, _>>()?, }, DataType::RunEndEncoded(_, _) => { let run_ends_child = arrays @@ -578,8 +610,8 @@ impl<'a> MutableArrayData<'a> { .map(|array| &array.child_data()[1]) .collect::>(); 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::>(); - MutableArrayData::new(child_arrays, use_nulls, array_capacity) + MutableArrayData::try_new(child_arrays, use_nulls, array_capacity) }) - .collect::>(), + .collect::, _>>()?, }; // 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 46f309556564..dec834a29dfd 100644 --- a/arrow-select/src/concat.rs +++ b/arrow-select/src/concat.rs @@ -579,7 +579,7 @@ pub fn concat(arrays: &[&dyn Array]) -> Result { fn concat_fallback(arrays: &[&dyn Array], capacity: Capacities) -> Result { let array_data: Vec<_> = arrays.iter().map(|a| a.to_data()).collect::>(); 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,51 @@ mod tests { ); } + #[test] + fn concat_string_view_dictionary_overflow_returns_err() { + // Two independently-built `Dictionary` 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::::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::::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::::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::::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 d84370947957..efd3493b64aa 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -768,7 +768,7 @@ fn interleave_fallback( ) -> Result { 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,58 @@ mod tests { ); } + #[test] + fn test_interleave_string_view_dictionary_overflow_returns_err() { + // Two independently-built `Dictionary` 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::::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::::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::::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::::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() {