This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11663-b05fbc7aede50f971e3b21f9254f7591feea5317 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit b858a1f951e5599ec639e9e27e0bd293fd304b6f Author: ump45nose <[email protected]> AuthorDate: Wed Oct 28 15:12:28 2026 +0000 fix(arrow-select): merge dictionaries in zip instead of panicking on key overflow (#11663) # Which issue does this PR close? - Closes #11594. # Rationale for this change `zip` on two dictionary arrays panicked with `DictionaryKeyOverflowError` when the two dictionaries together had more values than the key type can address, even when the values were the same in both and a merged dictionary would fit. Whether it panicked also depended on the mask, because long fragmented masks already went through `interleave`, which merges dictionaries. # What changes are included in this PR? `zip_impl` now builds its `MutableArrayData` with `try_new`. When that returns `DictionaryKeyOverflowError`, it builds the result with `interleave` instead, which merges the dictionaries (and returns an error rather than panicking if the merged values still don't fit). `interleave_arrays` now takes the scalar flags so this fallback also works when one or both sides are scalars. Inputs that didn't overflow take the same path as before. # Are these changes tested? Yes. `test_zip_dictionaries_whose_combined_values_overflow_the_key_type` zips two `Int8` dictionaries of the same 100 values for array/array, array/scalar, scalar/array and scalar/scalar inputs, and compares the decoded values and the data type with the expected ones. It panics without the fix. # Are there any significant user-facing changes? No. AI disclosure: I used an AI coding assistant to write the fix, the test and this description, and I reviewed all of it. The part worth a reviewer's look is the fallback in `zip_impl`, which only triggers on `DictionaryKeyOverflowError`. --- arrow-select/src/zip.rs | 74 +++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 69 insertions(+), 5 deletions(-) diff --git a/arrow-select/src/zip.rs b/arrow-select/src/zip.rs index a0381eff1c..911b3c6cf3 100644 --- a/arrow-select/src/zip.rs +++ b/arrow-select/src/zip.rs @@ -168,14 +168,24 @@ fn should_use_interleave(mask: &BooleanBuffer) -> bool { fn interleave_arrays( mask: &BooleanBuffer, truthy: &ArrayData, + truthy_is_scalar: bool, falsy: &ArrayData, + falsy_is_scalar: bool, ) -> Result<ArrayRef, ArrowError> { let truthy = make_array(truthy.clone()); let falsy = make_array(falsy.clone()); let indices: Vec<_> = mask .iter() .enumerate() - .map(|(idx, selected)| (usize::from(!selected), idx)) + .map(|(idx, selected)| { + let is_scalar = if selected { + truthy_is_scalar + } else { + falsy_is_scalar + }; + // A scalar is a single row array, so every row reads its first value + (usize::from(!selected), if is_scalar { 0 } else { idx }) + }) .collect(); crate::interleave::interleave(&[truthy.as_ref(), falsy.as_ref()], &indices) } @@ -189,10 +199,25 @@ fn zip_impl( ) -> Result<ArrayRef, ArrowError> { let mask_buffer = maybe_prep_null_mask_filter(mask); if !truthy_is_scalar && !falsy_is_scalar && should_use_interleave(&mask_buffer) { - return interleave_arrays(&mask_buffer, truthy, falsy); + return interleave_arrays(&mask_buffer, truthy, false, falsy, false); } - let mut mutable = MutableArrayData::new(vec![truthy, falsy], false, truthy.len()); + let mut mutable = match MutableArrayData::try_new(vec![truthy, falsy], false, truthy.len()) { + Ok(mutable) => mutable, + // `MutableArrayData` concatenates the two dictionaries, which can overflow + // the key type even when their merged values fit, so let `interleave` + // merge them instead + Err(ArrowError::DictionaryKeyOverflowError) => { + return interleave_arrays( + &mask_buffer, + truthy, + truthy_is_scalar, + falsy, + falsy_is_scalar, + ); + } + Err(e) => return Err(e), + }; // the SlicesIterator slices only the true values. So the gaps left by this iterator we need to // fill with falsy values @@ -923,7 +948,7 @@ impl<T: ByteViewType> ZipImpl for ByteViewScalarImpl<T> { #[cfg(test)] mod test { use super::*; - use arrow_array::types::Int32Type; + use arrow_array::types::{Int8Type, Int32Type}; #[test] fn test_count_true_runs() { @@ -951,6 +976,45 @@ mod test { assert_eq!(count_true_runs(&offset_mask.slice(3, 130)), 3); } + #[test] + fn test_zip_dictionaries_whose_combined_values_overflow_the_key_type() { + // Concatenating two dictionaries of 100 values overflows Int8 keys, but + // the values are the same in both, so the merged dictionary fits + let values: Vec<String> = (0..100).map(|i| format!("v{i}")).collect(); + let dictionary = |keys: Vec<i8>| -> ArrayRef { + Arc::new( + DictionaryArray::<Int8Type>::try_new( + Int8Array::from(keys), + Arc::new(StringArray::from(values.clone())), + ) + .unwrap(), + ) + }; + let mask = BooleanArray::from(vec![Some(true), Some(false), None, Some(true)]); + let truthy = dictionary(vec![0, 1, 2, 3]); + let falsy = dictionary(vec![10, 20, 30, 40]); + let truthy_scalar = Scalar::new(dictionary(vec![7])); + let falsy_scalar = Scalar::new(dictionary(vec![5])); + + let cases: [(&dyn Datum, &dyn Datum, [&str; 4]); 4] = [ + (&truthy, &falsy, ["v0", "v20", "v30", "v3"]), + (&truthy, &falsy_scalar, ["v0", "v5", "v5", "v3"]), + (&truthy_scalar, &falsy, ["v7", "v20", "v30", "v7"]), + (&truthy_scalar, &falsy_scalar, ["v7", "v5", "v5", "v7"]), + ]; + for (truthy_datum, falsy_datum, expected) in cases { + let actual = zip(&mask, truthy_datum, falsy_datum).unwrap(); + assert_eq!(actual.data_type(), truthy.data_type()); + let actual: Vec<_> = actual + .as_dictionary::<Int8Type>() + .downcast_dict::<StringArray>() + .unwrap() + .into_iter() + .collect(); + assert_eq!(actual, expected.map(Some)); + } + } + #[test] fn test_should_use_interleave() { let short: BooleanBuffer = (0..64).map(|i| i % 2 == 0).collect(); @@ -989,7 +1053,7 @@ mod test { let falsy = Int32Array::from(vec![Some(10), Some(20), None, Some(40)]).to_data(); let expected = Int32Array::from(vec![Some(1), Some(20), Some(3), Some(40)]); - let actual = interleave_arrays(&mask, &truthy, &falsy).unwrap(); + let actual = interleave_arrays(&mask, &truthy, false, &falsy, false).unwrap(); assert_eq!(actual.as_primitive::<Int32Type>(), &expected); }
