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-11925-9532701801cf2aada8b1450170c691192ff92bc7 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit 14c2fef335c4396552d60d69502950f2d0424c70 Author: NingLin Liu <[email protected]> AuthorDate: Wed Oct 28 15:11:25 2026 +0000 fix(arrow-select): merge dictionaries in merge and merge_n instead of panicking on key overflow (#11925) # Which issue does this PR close? - Closes #11834 # Rationale for this change `merge` and `merge_n` build their output with the infallible `MutableArrayData::new`, which concatenates the input dictionaries and panics with `DictionaryKeyOverflowError` when the combined values do not fit the key type, even when every input has the same values and a merged dictionary would fit. `concat` and `interleave` already merge dictionaries in that case. # What changes are included in this PR? - Both kernels now call `MutableArrayData::try_new`. On `DictionaryKeyOverflowError` they fall back to `interleave`, which merges the dictionaries (or returns an error if the merged dictionary still does not fit). Other errors are returned as before, and the non-dictionary fast path is unchanged. - `merge` builds the `(array, row)` pairs from the mask, treating null mask entries as false and reading row 0 for a scalar side, the same way the slice-based path does. - `merge_n` builds the pairs from the indices. `interleave` has no null index, so `None` rows read from an extra one-row null array of the same type. - `merge` with two scalars delegates to `zip`, which is handled by #11663, so it is not touched here. # Are these changes tested? Yes. Two new tests use two `Int8`-keyed dictionaries with the same 100 values, so concatenation overflows but the merged dictionary fits: - `test_merge_dictionaries_whose_combined_values_overflow_the_key_type`: array/array, array/scalar and scalar/array, with a null in the mask and a null key in the input. - `test_merge_n_dictionaries_whose_combined_values_overflow_the_key_type`: `Option<usize>` indices with `None` rows and a null key. Both panic on `main` and pass with this change. `cargo test -p arrow-select` and `cargo clippy -p arrow-select --all-targets -- -D warnings` pass. # Are there any user-facing changes? `merge` and `merge_n` return a merged dictionary instead of panicking. No API changes. This PR was prepared with AI assistance and reviewed/tested by the author. --- arrow-select/src/merge.rs | 166 ++++++++++++++++++++++++++++++++++++++++- arrow/benches/merge_kernels.rs | 33 ++++++++ 2 files changed, 195 insertions(+), 4 deletions(-) diff --git a/arrow-select/src/merge.rs b/arrow-select/src/merge.rs index be69054cbd..6aa84eaf54 100644 --- a/arrow-select/src/merge.rs +++ b/arrow-select/src/merge.rs @@ -19,7 +19,9 @@ use crate::filter::{SlicesIterator, prep_null_mask_filter}; use crate::zip::zip; -use arrow_array::{Array, ArrayRef, BooleanArray, Datum, make_array, new_empty_array}; +use arrow_array::{ + Array, ArrayRef, BooleanArray, Datum, make_array, new_empty_array, new_null_array, +}; use arrow_data::ArrayData; use arrow_data::transform::MutableArrayData; use arrow_schema::ArrowError; @@ -144,7 +146,16 @@ pub fn merge_n(values: &[&dyn Array], indices: &[impl MergeIndex]) -> Result<Arr let data: Vec<ArrayData> = values.iter().map(|a| a.to_data()).collect(); let data_refs = data.iter().collect(); - let mut mutable = MutableArrayData::new(data_refs, true, indices.len()); + let mut mutable = match MutableArrayData::try_new(data_refs, true, indices.len()) { + Ok(mutable) => mutable, + // `MutableArrayData` concatenates the dictionaries, which can overflow + // the key type even when their merged values fit, so let `interleave` + // merge them instead + Err(ArrowError::DictionaryKeyOverflowError) => { + return merge_n_with_interleave(values, indices); + } + Err(e) => return Err(e), + }; // This loop extends the mutable array by taking slices from the partial results. // @@ -182,6 +193,61 @@ pub fn merge_n(values: &[&dyn Array], indices: &[impl MergeIndex]) -> Result<Arr Ok(make_array(mutable.freeze())) } +/// [merge_n] through [interleave](crate::interleave::interleave), which merges +/// dictionaries instead of concatenating them +fn merge_n_with_interleave( + values: &[&dyn Array], + indices: &[impl MergeIndex], +) -> Result<ArrayRef, ArrowError> { + // `interleave` has no null index, so null rows read from an extra + // single null array + let nulls = new_null_array(values[0].data_type(), 1); + let mut arrays = values.to_vec(); + arrays.push(nulls.as_ref()); + + let mut take_offsets = vec![0; values.len()]; + let interleave_indices: Vec<_> = indices + .iter() + .map(|ix| match ix.index() { + None => (values.len(), 0), + Some(index) => { + let offset = take_offsets[index]; + take_offsets[index] += 1; + (index, offset) + } + }) + .collect(); + crate::interleave::interleave(&arrays, &interleave_indices) +} + +/// [merge] through [interleave](crate::interleave::interleave), which merges +/// dictionaries instead of concatenating them +fn merge_with_interleave( + mask: &BooleanArray, + truthy: &dyn Array, + truthy_is_scalar: bool, + falsy: &dyn Array, + falsy_is_scalar: bool, +) -> Result<ArrayRef, ArrowError> { + let mut truthy_offset = 0; + let mut falsy_offset = 0; + let indices: Vec<_> = (0..mask.len()) + .map(|idx| { + // Nulls in the mask are treated as false + if mask.is_valid(idx) && mask.value(idx) { + let offset = if truthy_is_scalar { 0 } else { truthy_offset }; + truthy_offset += 1; + (0, offset) + } else { + let offset = if falsy_is_scalar { 0 } else { falsy_offset }; + falsy_offset += 1; + (1, offset) + } + }) + .collect(); + crate::interleave::interleave(&[truthy, falsy], &indices) +} + /// Merges two arrays in the order specified by a boolean mask. /// /// This algorithm is a variant of [zip] that does not require the truthy and @@ -244,7 +310,20 @@ pub fn merge( let falsy = falsy_array.to_data(); let truthy = truthy_array.to_data(); - let mut mutable = MutableArrayData::new(vec![&truthy, &falsy], false, mask.len()); + let mut mutable = match MutableArrayData::try_new(vec![&truthy, &falsy], false, mask.len()) { + Ok(mutable) => mutable, + // See `merge_n`: let `interleave` merge the dictionaries + Err(ArrowError::DictionaryKeyOverflowError) => { + return merge_with_interleave( + mask, + truthy_array, + truthy_is_scalar, + falsy_array, + 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 @@ -311,8 +390,13 @@ pub fn merge( mod tests { use crate::merge::{MergeIndex, merge, merge_n}; use arrow_array::cast::AsArray; - use arrow_array::{Array, BooleanArray, Datum, Int32Array, Scalar, StringArray, UInt64Array}; + use arrow_array::types::Int8Type; + use arrow_array::{ + Array, ArrayRef, BooleanArray, Datum, DictionaryArray, Int8Array, Int32Array, Scalar, + StringArray, UInt64Array, + }; use arrow_schema::ArrowError::InvalidArgumentError; + use std::sync::Arc; #[derive(PartialEq, Eq, Copy, Clone)] struct CompactMergeIndex { @@ -615,4 +699,78 @@ mod tests { assert!(matches!(merged, Err(InvalidArgumentError { .. }))); } + + /// Two dictionaries of 100 values, the same in both. Concatenating them + /// overflows Int8 keys, but the merged dictionary fits + fn overflowing_dictionary(keys: Vec<Option<i8>>) -> ArrayRef { + let values: Vec<String> = (0..100).map(|i| format!("v{i}")).collect(); + Arc::new( + DictionaryArray::<Int8Type>::try_new( + Int8Array::from(keys), + Arc::new(StringArray::from(values)), + ) + .unwrap(), + ) + } + + fn dictionary_strings(array: &dyn Array) -> Vec<Option<&str>> { + array + .as_dictionary::<Int8Type>() + .downcast_dict::<StringArray>() + .unwrap() + .into_iter() + .collect() + } + + #[test] + fn test_merge_dictionaries_whose_combined_values_overflow_the_key_type() { + let mask = BooleanArray::from(vec![Some(true), Some(false), None, Some(true), Some(false)]); + let truthy = overflowing_dictionary(vec![Some(0), None, Some(2)]); + let falsy = overflowing_dictionary(vec![Some(10), Some(20), Some(30)]); + let truthy_scalar = Scalar::new(overflowing_dictionary(vec![Some(7)])); + let falsy_scalar = Scalar::new(overflowing_dictionary(vec![Some(5)])); + + let check = + |truthy_datum: &dyn Datum, falsy_datum: &dyn Datum, expected: [Option<&str>; 5]| { + let merged = merge(&mask, truthy_datum, falsy_datum).unwrap(); + assert_eq!(merged.data_type(), truthy.data_type()); + assert_eq!(dictionary_strings(merged.as_ref()), expected); + }; + check( + &truthy, + &falsy, + [Some("v0"), Some("v10"), Some("v20"), None, Some("v30")], + ); + check( + &truthy, + &falsy_scalar, + [Some("v0"), Some("v5"), Some("v5"), None, Some("v5")], + ); + check( + &truthy_scalar, + &falsy, + [ + Some("v7"), + Some("v10"), + Some("v20"), + Some("v7"), + Some("v30"), + ], + ); + } + + #[test] + fn test_merge_n_dictionaries_whose_combined_values_overflow_the_key_type() { + let a = overflowing_dictionary(vec![Some(0), Some(1)]); + let b = overflowing_dictionary(vec![Some(50), None]); + let indices = vec![Some(1), None, Some(0), Some(0), Some(1), None]; + + let merged = merge_n(&[a.as_ref(), b.as_ref()], &indices).unwrap(); + + assert_eq!(merged.data_type(), a.data_type()); + assert_eq!( + dictionary_strings(merged.as_ref()), + [Some("v50"), None, Some("v0"), Some("v1"), None, None] + ); + } } diff --git a/arrow/benches/merge_kernels.rs b/arrow/benches/merge_kernels.rs index 2744e6618f..08f541bfdd 100644 --- a/arrow/benches/merge_kernels.rs +++ b/arrow/benches/merge_kernels.rs @@ -225,6 +225,36 @@ fn bench_merge_input_on_all_masks( } } +/// An Int8-keyed dictionary of `len` random keys into `distinct` string values. +/// Each call builds its own values array, so two calls share values but not +/// the dictionary itself +fn generate_int8_dictionary(seed: u64, len: usize, distinct: usize) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(seed); + let values = StringArray::from_iter_values((0..distinct).map(|i| format!("value_{i}"))); + let keys = Int8Array::from_iter_values((0..len).map(|_| rng.random_range(0..distinct) as i8)); + Arc::new(DictionaryArray::<Int8Type>::try_new(keys, Arc::new(values)).unwrap()) +} + +fn bench_merge_dictionaries(c: &mut Criterion) { + const ARRAY_LEN: usize = 8192; + + let mut group = c.benchmark_group(format!("merge_{ARRAY_LEN}_from_int8_dictionaries")); + let masks = mask_cases(ARRAY_LEN); + + // 2 x 50 values fit in Int8 keys when the dictionaries are concatenated, + // 2 x 100 values do not, so merging them has to merge the dictionaries + for (description, distinct) in [ + ("array_vs_array_concatenated_values_fit", 50), + ("array_vs_array_concatenated_values_overflow", 100), + ] { + let truthy = generate_int8_dictionary(42, ARRAY_LEN, distinct); + let falsy = generate_int8_dictionary(18, ARRAY_LEN, distinct); + bench_merge_input_on_all_masks(description, &mut group, &masks, &truthy, &falsy); + } + + group.finish(); +} + fn add_benchmark(c: &mut Criterion) { // Primitive bench_merge_on_input_generator( @@ -274,6 +304,9 @@ fn add_benchmark(c: &mut Criterion) { _marker: std::marker::PhantomData, }, ); + + // Dictionaries + bench_merge_dictionaries(c); } criterion_group!(benches, add_benchmark);
