arrow_select/
dictionary.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Dictionary utilities for Arrow arrays
6
7use std::{collections::HashMap, sync::Arc};
8
9use crate::filter::filter;
10use crate::interleave::interleave;
11use ahash::RandomState;
12use arrow_array::builder::BooleanBufferBuilder;
13use arrow_array::types::{
14    ArrowDictionaryKeyType, ArrowPrimitiveType, BinaryType, ByteArrayType, LargeBinaryType,
15    LargeUtf8Type, Utf8Type,
16};
17use arrow_array::{
18    AnyDictionaryArray, Array, ArrayRef, ArrowNativeTypeOp, BooleanArray, DictionaryArray,
19    GenericByteArray, PrimitiveArray, downcast_dictionary_array,
20};
21use arrow_array::{cast::AsArray, downcast_primitive};
22use arrow_buffer::{ArrowNativeType, BooleanBuffer, ScalarBuffer, ToByteSlice};
23use arrow_schema::{ArrowError, DataType};
24
25/// Garbage collects a [DictionaryArray] by removing unreferenced values.
26///
27/// Returns a new [DictionaryArray] such that there are no values
28/// that are not referenced by at least one key. There may still be duplicate
29/// values.
30///
31/// See also [`garbage_collect_any_dictionary`] if you need to handle multiple dictionary types
32pub fn garbage_collect_dictionary<K: ArrowDictionaryKeyType>(
33    dictionary: &DictionaryArray<K>,
34) -> Result<DictionaryArray<K>, ArrowError> {
35    let keys = dictionary.keys();
36    let values = dictionary.values();
37
38    let mask = dictionary.occupancy();
39
40    // If no work to do, return the original dictionary
41    if mask.count_set_bits() == values.len() {
42        return Ok(dictionary.clone());
43    }
44
45    // Create a mapping from the old keys to the new keys, use a Vec for easy indexing
46    let mut key_remap = vec![K::Native::ZERO; values.len()];
47    for (new_idx, old_idx) in mask.set_indices().enumerate() {
48        key_remap[old_idx] =
49            K::Native::from_usize(new_idx).ok_or(ArrowError::DictionaryKeyOverflowError)?;
50    }
51
52    // ... and then build the new keys array
53    let new_keys = keys.unary(|key| {
54        key_remap
55            .get(key.as_usize())
56            .copied()
57            // nulls may be present in the keys, and they will have arbitrary value; we don't care
58            // and can safely return zero
59            .unwrap_or(K::Native::ZERO)
60    });
61
62    // Create a new values array by filtering using the mask
63    let values = filter(dictionary.values(), &BooleanArray::new(mask, None))?;
64
65    DictionaryArray::try_new(new_keys, values)
66}
67
68/// Equivalent to [`garbage_collect_dictionary`] but without requiring casting to a specific key type.
69pub fn garbage_collect_any_dictionary(
70    dictionary: &dyn AnyDictionaryArray,
71) -> Result<ArrayRef, ArrowError> {
72    let dictionary: &dyn Array = dictionary;
73    downcast_dictionary_array!(
74        dictionary => garbage_collect_dictionary(dictionary).map(|dict| Arc::new(dict) as ArrayRef),
75        _ => unreachable!("have a dictionary array")
76    )
77}
78
79/// A best effort interner that maintains a fixed number of buckets
80/// and interns keys based on their hash value
81///
82/// Hash collisions will result in replacement
83struct Interner<'a, V> {
84    state: RandomState,
85    buckets: Vec<Option<InternerBucket<'a, V>>>,
86    shift: u32,
87}
88
89/// A single bucket in [`Interner`].
90type InternerBucket<'a, V> = (Option<&'a [u8]>, V);
91
92impl<'a, V> Interner<'a, V> {
93    /// Capacity controls the number of unique buckets allocated within the Interner
94    ///
95    /// A larger capacity reduces the probability of hash collisions, and should be set
96    /// based on an approximation of the upper bound of unique values
97    fn new(capacity: usize) -> Self {
98        // Add additional buckets to help reduce collisions
99        let shift = (capacity as u64 + 128).leading_zeros();
100        let num_buckets = (u64::MAX >> shift) as usize;
101        let buckets = (0..num_buckets.saturating_add(1)).map(|_| None).collect();
102        Self {
103            // A fixed seed to ensure deterministic behaviour
104            state: RandomState::with_seeds(0, 0, 0, 0),
105            buckets,
106            shift,
107        }
108    }
109
110    fn intern<F: FnOnce() -> Result<V, E>, E>(
111        &mut self,
112        new: Option<&'a [u8]>,
113        f: F,
114    ) -> Result<&V, E> {
115        let hash = self.state.hash_one(new);
116        let bucket_idx = hash >> self.shift;
117        Ok(match &mut self.buckets[bucket_idx as usize] {
118            Some((current, v)) => {
119                if *current != new {
120                    *v = f()?;
121                    *current = new;
122                }
123                v
124            }
125            slot => &slot.insert((new, f()?)).1,
126        })
127    }
128}
129
130// Sparse hashing pays off only when retained values greatly outnumber key lookups.
131pub(crate) const MAX_DENSE_MAPPING_VALUES_PER_ROW: usize = 256;
132
133pub(crate) enum KeyMapping<K: ArrowDictionaryKeyType> {
134    Dense(Vec<K::Native>),
135    Sparse(HashMap<usize, K::Native, RandomState>),
136}
137
138impl<K: ArrowDictionaryKeyType> KeyMapping<K> {
139    fn new_dense(len: usize) -> Self {
140        Self::Dense(vec![K::Native::ZERO; len])
141    }
142
143    fn new_sparse(capacity: usize) -> Self {
144        Self::Sparse(HashMap::with_capacity_and_hasher(
145            capacity,
146            RandomState::with_seeds(0, 0, 0, 0),
147        ))
148    }
149
150    fn insert(&mut self, old_key: usize, new_key: K::Native) {
151        match self {
152            Self::Dense(mapping) => mapping[old_key] = new_key,
153            Self::Sparse(mapping) => {
154                mapping.insert(old_key, new_key);
155            }
156        }
157    }
158
159    #[inline]
160    pub fn get(&self, old_key: usize) -> Option<K::Native> {
161        match self {
162            Self::Dense(mapping) => mapping.get(old_key).copied(),
163            Self::Sparse(mapping) => mapping.get(&old_key).copied(),
164        }
165    }
166}
167
168pub(crate) struct MergedDictionaries<K: ArrowDictionaryKeyType> {
169    /// Maps each input dictionary key to its merged key.
170    pub key_mappings: Vec<KeyMapping<K>>,
171    /// The new values
172    pub values: ArrayRef,
173}
174
175/// Performs a cheap, pointer-based comparison of two byte array
176///
177/// See [`ScalarBuffer::ptr_eq`]
178fn bytes_ptr_eq<T: ByteArrayType>(a: &dyn Array, b: &dyn Array) -> bool {
179    match (a.as_bytes_opt::<T>(), b.as_bytes_opt::<T>()) {
180        (Some(a), Some(b)) => {
181            let values_eq = a.values().ptr_eq(b.values()) && a.offsets().ptr_eq(b.offsets());
182            match (a.nulls(), b.nulls()) {
183                (Some(a), Some(b)) => values_eq && a.inner().ptr_eq(b.inner()),
184                (None, None) => values_eq,
185                _ => false,
186            }
187        }
188        _ => false,
189    }
190}
191
192/// Returns `true` if two child value arrays have the same layout and backing buffers.
193///
194/// This first checks whether the [`ArrayRef`]s point to the same allocation, then uses
195/// direct buffer comparisons for byte arrays and [`arrow_data::ArrayData::ptr_eq`] for
196/// other types.
197///
198/// This checks storage identity, not logical equality. Arrays with equal values in
199/// different buffers return `false`. Arrays with no allocated storage, such as two empty
200/// primitive arrays built from empty `Vec`s or two [`arrow_array::NullArray`]s of the same
201/// length, return `true`.
202pub(crate) fn child_values_ptr_eq(left: &ArrayRef, right: &ArrayRef) -> bool {
203    if Arc::ptr_eq(left, right) {
204        return true;
205    }
206
207    match left.data_type() {
208        DataType::Utf8 => bytes_ptr_eq::<Utf8Type>(left.as_ref(), right.as_ref()),
209        DataType::LargeUtf8 => bytes_ptr_eq::<LargeUtf8Type>(left.as_ref(), right.as_ref()),
210        DataType::Binary => bytes_ptr_eq::<BinaryType>(left.as_ref(), right.as_ref()),
211        DataType::LargeBinary => bytes_ptr_eq::<LargeBinaryType>(left.as_ref(), right.as_ref()),
212        _ => left.to_data().ptr_eq(&right.to_data()),
213    }
214}
215
216/// A type-erased function that compares two array for pointer equality
217type PtrEq = fn(&dyn Array, &dyn Array) -> bool;
218
219/// A weak heuristic of whether to merge dictionary values that aims to only
220/// perform the expensive merge computation when it is likely to yield at least
221/// some return over the naive approach used by MutableArrayData
222///
223/// `len` is the total length of the merged output
224///
225/// Returns `(should_merge, has_overflow)` where:
226/// - `should_merge`: whether dictionary values should be merged
227/// - `has_overflow`: whether the combined dictionary values would overflow the key type
228pub(crate) fn should_merge_dictionary_values<K: ArrowDictionaryKeyType>(
229    dictionaries: &[&DictionaryArray<K>],
230    len: usize,
231) -> (bool, bool) {
232    use DataType::*;
233    let first_values = dictionaries[0].values().as_ref();
234    let ptr_eq: PtrEq = match first_values.data_type() {
235        Utf8 => bytes_ptr_eq::<Utf8Type>,
236        LargeUtf8 => bytes_ptr_eq::<LargeUtf8Type>,
237        Binary => bytes_ptr_eq::<BinaryType>,
238        LargeBinary => bytes_ptr_eq::<LargeBinaryType>,
239        dt => {
240            if !dt.is_primitive() {
241                return (
242                    false,
243                    K::Native::from_usize(dictionaries.iter().map(|d| d.values().len()).sum())
244                        .is_none(),
245                );
246            }
247            |a, b| a.to_data().ptr_eq(&b.to_data())
248        }
249    };
250
251    let mut single_dictionary = true;
252    let mut total_values = first_values.len();
253    for dict in dictionaries.iter().skip(1) {
254        let values = dict.values().as_ref();
255        total_values += values.len();
256        if single_dictionary {
257            single_dictionary = ptr_eq(first_values, values)
258        }
259    }
260
261    let overflow = K::Native::from_usize(total_values).is_none();
262    let values_exceed_length = total_values >= len;
263
264    (
265        !single_dictionary && (overflow || values_exceed_length),
266        overflow,
267    )
268}
269
270/// Given an array of dictionaries and an optional key mask compute a values array
271/// containing referenced values, along with mappings from the [`DictionaryArray`]
272/// keys to the new keys within this values array. Best-effort will be made to ensure
273/// that the dictionary values are unique
274///
275/// This method is meant to be very fast and the output dictionary values
276/// may not be unique, unlike `GenericByteDictionaryBuilder` which is slower
277/// but produces unique values. `key_lookups` is the number of times each
278/// input's keys will be remapped, and defaults to the input length.
279pub(crate) fn merge_dictionary_values<K: ArrowDictionaryKeyType>(
280    dictionaries: &[&DictionaryArray<K>],
281    masks: Option<&[BooleanBuffer]>,
282    key_lookups: Option<&[usize]>,
283) -> Result<MergedDictionaries<K>, ArrowError> {
284    let mut num_values = 0;
285
286    let mut values_arrays = Vec::with_capacity(dictionaries.len());
287    let mut value_slices = Vec::with_capacity(dictionaries.len());
288
289    for (idx, dictionary) in dictionaries.iter().enumerate() {
290        let mask = masks.and_then(|m| m.get(idx));
291        let key_mask_owned;
292        let key_mask = match (dictionary.nulls(), mask) {
293            (Some(n), None) => Some(n.inner()),
294            (None, Some(n)) => Some(n),
295            (Some(n), Some(m)) => {
296                key_mask_owned = n.inner() & m;
297                Some(&key_mask_owned)
298            }
299            (None, None) => None,
300        };
301        let keys = dictionary.keys().values();
302        let values = dictionary.values().as_ref();
303        let lookups = key_lookups.map_or(keys.len(), |lookups| lookups[idx]);
304        let use_sparse_mapping =
305            values.len() > lookups.saturating_mul(MAX_DENSE_MAPPING_VALUES_PER_ROW);
306
307        let selected_values = if use_sparse_mapping {
308            let referenced_keys = compute_referenced_keys(keys, key_mask);
309            let len = referenced_keys.len();
310            get_selected_values(values, referenced_keys.into_iter(), len)
311        } else {
312            let values_mask = compute_values_mask(keys, key_mask, values.len());
313            let len = values_mask.count_set_bits();
314            get_selected_values(values, values_mask.set_indices(), len)
315        };
316        num_values += selected_values.len();
317        value_slices.push((selected_values, use_sparse_mapping));
318        values_arrays.push(values)
319    }
320
321    // Map from value to new index
322    let mut interner = Interner::new(num_values);
323    // Interleave indices for new values array
324    let mut indices = Vec::with_capacity(num_values);
325
326    // Compute the mapping for each dictionary
327    let key_mappings = dictionaries
328        .iter()
329        .enumerate()
330        .zip(value_slices)
331        .map(
332            |((dictionary_idx, dictionary), (values, use_sparse_mapping))| {
333                let mut mapping = match use_sparse_mapping {
334                    true => KeyMapping::new_sparse(values.len()),
335                    false => KeyMapping::new_dense(dictionary.values().len()),
336                };
337
338                for (value_idx, value) in values {
339                    let new_key =
340                        *interner.intern(value, || match K::Native::from_usize(indices.len()) {
341                            Some(idx) => {
342                                indices.push((dictionary_idx, value_idx));
343                                Ok(idx)
344                            }
345                            None => Err(ArrowError::DictionaryKeyOverflowError),
346                        })?;
347                    mapping.insert(value_idx, new_key);
348                }
349                Ok(mapping)
350            },
351        )
352        .collect::<Result<Vec<_>, ArrowError>>()?;
353
354    Ok(MergedDictionaries {
355        key_mappings,
356        values: interleave(&values_arrays, &indices)?,
357    })
358}
359
360/// Merges referenced entries from multiple dictionary values arrays.
361///
362/// `values_masks[i]` corresponds to `values_arrays[i]`, with each set bit identifying a
363/// value to retain. Returns the merged values and a mapping from each source value index
364/// to its index in the merged array. Mapping entries for unreferenced values are zero and
365/// must not be used.
366///
367/// Referenced values are gathered without hashing when their count fits in `K`'s key space,
368/// preserving duplicates across sources. Otherwise, values are deduplicated on a best-effort
369/// basis to avoid key overflow.
370pub(crate) fn merge_dictionary_values_with_masks<K: ArrowDictionaryKeyType>(
371    values_arrays: &[&dyn Array],
372    values_masks: &[BooleanBuffer],
373) -> Result<MergedDictionaries<K>, ArrowError> {
374    assert_eq!(values_arrays.len(), values_masks.len());
375    let num_values = values_masks.iter().map(BooleanBuffer::count_set_bits).sum();
376
377    // Dictionary values need not be unique. Avoid hashing when the referenced value count
378    // fits in `K`'s key space, and only deduplicate to avoid overflow.
379    if num_values == 0 || K::Native::from_usize(num_values - 1).is_some() {
380        let zero = K::Native::from_usize(0).unwrap();
381        let mut indices = Vec::with_capacity(num_values);
382        let key_mappings = values_arrays
383            .iter()
384            .zip(values_masks)
385            .enumerate()
386            .map(|(dictionary_idx, (values, mask))| {
387                let mut mapping = vec![zero; values.len()];
388                for value_idx in mask.set_indices() {
389                    mapping[value_idx] = K::Native::from_usize(indices.len()).unwrap();
390                    indices.push((dictionary_idx, value_idx));
391                }
392                KeyMapping::Dense(mapping)
393            })
394            .collect();
395        return Ok(MergedDictionaries {
396            key_mappings,
397            values: interleave(values_arrays, &indices)?,
398        });
399    }
400
401    let value_slices = values_arrays
402        .iter()
403        .zip(values_masks)
404        .map(|(values, mask)| {
405            let len = mask.count_set_bits();
406            get_selected_values(*values, mask.set_indices(), len)
407        })
408        .collect::<Vec<_>>();
409    let mut interner = Interner::new(num_values);
410    let mut indices = Vec::with_capacity(num_values);
411
412    let key_mappings = values_arrays
413        .iter()
414        .enumerate()
415        .zip(value_slices)
416        .map(|((dictionary_idx, value_array), values)| {
417            let zero = K::Native::ZERO;
418            let mut mapping = vec![zero; value_array.len()];
419
420            for (value_idx, value) in values {
421                mapping[value_idx] =
422                    *interner.intern(value, || match K::Native::from_usize(indices.len()) {
423                        Some(idx) => {
424                            indices.push((dictionary_idx, value_idx));
425                            Ok(idx)
426                        }
427                        None => Err(ArrowError::DictionaryKeyOverflowError),
428                    })?;
429            }
430            Ok(KeyMapping::Dense(mapping))
431        })
432        .collect::<Result<Vec<_>, ArrowError>>()?;
433
434    Ok(MergedDictionaries {
435        key_mappings,
436        values: interleave(values_arrays, &indices)?,
437    })
438}
439
440/// Returns the unique dictionary keys referenced by the selected rows.
441///
442/// Sorting makes repeated keys adjacent for deduplication and preserves the
443/// ascending dictionary-key order of the dense mask path. This keeps the sparse
444/// and dense paths' merged values ordered consistently.
445fn compute_referenced_keys<K: ArrowNativeType>(
446    keys: &ScalarBuffer<K>,
447    mask: Option<&BooleanBuffer>,
448) -> Vec<usize> {
449    let mut indices = match mask {
450        Some(mask) => mask
451            .set_indices()
452            .map(|idx| keys[idx].as_usize())
453            .collect::<Vec<_>>(),
454        None => keys.iter().map(|key| key.as_usize()).collect(),
455    };
456    indices.sort_unstable();
457    indices.dedup();
458    indices
459}
460
461/// Return a mask identifying the values that are referenced by keys in `dictionary`
462/// at the positions indicated by `selection`
463fn compute_values_mask<K: ArrowNativeType>(
464    keys: &ScalarBuffer<K>,
465    mask: Option<&BooleanBuffer>,
466    max_key: usize,
467) -> BooleanBuffer {
468    let mut builder = BooleanBufferBuilder::new(max_key);
469    builder.advance(max_key);
470
471    match mask {
472        Some(n) => n
473            .set_indices()
474            .for_each(|idx| builder.set_bit(keys[idx].as_usize(), true)),
475        None => keys
476            .iter()
477            .for_each(|k| builder.set_bit(k.as_usize(), true)),
478    }
479    builder.finish()
480}
481
482fn selected_primitives_to_bytes<T: ArrowPrimitiveType, I: Iterator<Item = usize>>(
483    array: &PrimitiveArray<T>,
484    indices: I,
485    len: usize,
486) -> Vec<(usize, Option<&[u8]>)>
487where
488    T::Native: ToByteSlice,
489{
490    let mut out = Vec::with_capacity(len);
491    let values = array.values();
492    for idx in indices {
493        out.push((
494            idx,
495            array.is_valid(idx).then_some(values[idx].to_byte_slice()),
496        ))
497    }
498    out
499}
500
501macro_rules! selected_primitive_to_bytes_helper {
502    ($t:ty, $array:expr, $indices:expr, $len:expr) => {
503        selected_primitives_to_bytes::<$t, _>($array.as_primitive(), $indices, $len)
504    };
505}
506
507fn get_selected_values<I: Iterator<Item = usize>>(
508    array: &dyn Array,
509    indices: I,
510    len: usize,
511) -> Vec<(usize, Option<&[u8]>)> {
512    downcast_primitive! {
513        array.data_type() => (selected_primitive_to_bytes_helper, array, indices, len),
514        DataType::Utf8 => selected_bytes(array.as_string::<i32>(), indices, len),
515        DataType::LargeUtf8 => selected_bytes(array.as_string::<i64>(), indices, len),
516        DataType::Binary => selected_bytes(array.as_binary::<i32>(), indices, len),
517        DataType::LargeBinary => selected_bytes(array.as_binary::<i64>(), indices, len),
518        _ => unimplemented!("Dictionary merging for type {} is not implemented", array.data_type()),
519    }
520}
521
522fn selected_bytes<T: ByteArrayType, I: Iterator<Item = usize>>(
523    array: &GenericByteArray<T>,
524    indices: I,
525    len: usize,
526) -> Vec<(usize, Option<&[u8]>)> {
527    let mut out = Vec::with_capacity(len);
528    for idx in indices {
529        out.push((
530            idx,
531            array.is_valid(idx).then_some(array.value(idx).as_ref()),
532        ))
533    }
534    out
535}
536
537#[cfg(test)]
538mod tests {
539    use super::*;
540
541    use arrow_array::cast::as_string_array;
542    use arrow_array::types::{Int8Type, Int32Type, UInt8Type};
543    use arrow_array::{
544        ArrayRef, DictionaryArray, Int8Array, Int32Array, ListArray, NullArray, RunArray,
545        StringArray, StringViewArray, UInt8Array,
546    };
547    use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer, OffsetBuffer};
548    use arrow_schema::Field;
549    use std::sync::Arc;
550
551    fn assert_mapping(mapping: &KeyMapping<Int32Type>, expected: &[i32]) {
552        let KeyMapping::Dense(actual) = mapping else {
553            panic!("expected a dense mapping");
554        };
555        assert_eq!(actual, expected);
556    }
557
558    #[test]
559    fn test_child_values_ptr_eq_accepts_reusable_value_mapping() {
560        let dictionary = DictionaryArray::<Int8Type>::new(
561            Int8Array::from(vec![0, 1]),
562            Arc::new(StringArray::from(vec!["a", "b"])),
563        );
564        let same_values = Arc::clone(dictionary.values());
565
566        assert!(Arc::ptr_eq(dictionary.values(), &same_values));
567        assert!(child_values_ptr_eq(dictionary.values(), &same_values));
568
569        let shared_values = StringArray::from(vec!["a", "b"]);
570        let first_dictionary = DictionaryArray::<Int8Type>::new(
571            Int8Array::from(vec![0, 1]),
572            Arc::new(shared_values.clone()),
573        );
574        let cloned_dictionary =
575            DictionaryArray::<Int8Type>::new(Int8Array::from(vec![0, 1]), Arc::new(shared_values));
576
577        assert!(!Arc::ptr_eq(
578            first_dictionary.values(),
579            cloned_dictionary.values()
580        ));
581        assert!(child_values_ptr_eq(
582            first_dictionary.values(),
583            cloned_dictionary.values()
584        ));
585
586        // Run boundaries do not affect whether the child value mapping can be reused.
587        let run_values = StringArray::from(vec!["a", "b"]);
588        let first_runs =
589            RunArray::<Int32Type>::try_new(&Int32Array::from(vec![1, 3]), &run_values).unwrap();
590        let second_runs =
591            RunArray::<Int32Type>::try_new(&Int32Array::from(vec![2, 3]), &run_values).unwrap();
592
593        assert!(!Arc::ptr_eq(first_runs.values(), second_runs.values()));
594        assert!(child_values_ptr_eq(
595            first_runs.values(),
596            second_runs.values()
597        ));
598
599        // Arrays without allocated buffers still have a reusable value mapping.
600        let empty_left: ArrayRef = Arc::new(Int32Array::from(Vec::<i32>::new()));
601        let empty_right: ArrayRef = Arc::new(Int32Array::from(Vec::<i32>::new()));
602        assert!(child_values_ptr_eq(&empty_left, &empty_right));
603
604        let null_left: ArrayRef = Arc::new(NullArray::new(3));
605        let null_right: ArrayRef = Arc::new(NullArray::new(3));
606        assert!(child_values_ptr_eq(&null_left, &null_right));
607    }
608
609    #[test]
610    fn test_child_values_ptr_eq_rejects_different_storage_or_layout() {
611        let independent_left: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
612        let independent_right: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
613
614        assert!(!child_values_ptr_eq(&independent_left, &independent_right));
615
616        let storage = Int32Array::from(vec![1, 2, 3]);
617        let prefix: ArrayRef = Arc::new(storage.slice(0, 2));
618        let shifted: ArrayRef = Arc::new(storage.slice(1, 2));
619        let shorter: ArrayRef = Arc::new(storage.slice(0, 1));
620
621        assert!(!child_values_ptr_eq(&prefix, &shifted));
622        assert!(!child_values_ptr_eq(&prefix, &shorter));
623    }
624
625    #[test]
626    fn test_child_values_ptr_eq_rejects_different_nested_or_view_storage() {
627        // Shared parent metadata is insufficient when child or payload storage differs.
628        let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
629        let shared_offsets = OffsetBuffer::new(vec![0_i32, 2].into());
630        let first_list: ArrayRef = Arc::new(ListArray::new(
631            Arc::clone(&list_field),
632            shared_offsets.clone(),
633            Arc::new(Int32Array::from(vec![1, 2])),
634            None,
635        ));
636        let second_list: ArrayRef = Arc::new(ListArray::new(
637            list_field,
638            shared_offsets,
639            Arc::new(Int32Array::from(vec![1, 3])),
640            None,
641        ));
642        assert!(!child_values_ptr_eq(&first_list, &second_list));
643
644        let first_payload: &[u8] = b"same payload one";
645        let second_payload: &[u8] = b"same payload two";
646        let view = arrow_data::ByteView::new(first_payload.len() as u32, &first_payload[..4]);
647        let shared_views = ScalarBuffer::from(vec![view.into()]);
648        let first_view: ArrayRef = Arc::new(StringViewArray::new(
649            shared_views.clone(),
650            vec![Buffer::from_slice_ref(first_payload)],
651            None,
652        ));
653        let second_view: ArrayRef = Arc::new(StringViewArray::new(
654            shared_views,
655            vec![Buffer::from_slice_ref(second_payload)],
656            None,
657        ));
658        assert!(!child_values_ptr_eq(&first_view, &second_view));
659    }
660
661    #[test]
662    fn test_garbage_collect_i32_dictionary() {
663        let values = StringArray::from_iter_values(["a", "b", "c", "d"]);
664        let keys = Int32Array::from_iter_values([0, 1, 1, 3, 0, 0, 1]);
665        let dict = DictionaryArray::<Int32Type>::new(keys, Arc::new(values));
666
667        // Only "a", "b", "d" are referenced, "c" is not
668        let gc = garbage_collect_dictionary(&dict).unwrap();
669
670        let expected_values = StringArray::from_iter_values(["a", "b", "d"]);
671        let expected_keys = Int32Array::from_iter_values([0, 1, 1, 2, 0, 0, 1]);
672        let expected = DictionaryArray::<Int32Type>::new(expected_keys, Arc::new(expected_values));
673
674        assert_eq!(gc, expected);
675    }
676
677    #[test]
678    fn test_garbage_collect_any_dictionary() {
679        let values = StringArray::from_iter_values(["a", "b", "c", "d"]);
680        let keys = Int32Array::from_iter_values([0, 1, 1, 3, 0, 0, 1]);
681        let dict = DictionaryArray::<Int32Type>::new(keys, Arc::new(values));
682
683        let gc = garbage_collect_any_dictionary(&dict).unwrap();
684
685        let expected_values = StringArray::from_iter_values(["a", "b", "d"]);
686        let expected_keys = Int32Array::from_iter_values([0, 1, 1, 2, 0, 0, 1]);
687        let expected = DictionaryArray::<Int32Type>::new(expected_keys, Arc::new(expected_values));
688
689        assert_eq!(gc.as_ref(), &expected);
690    }
691
692    #[test]
693    fn test_garbage_collect_with_nulls() {
694        let values = StringArray::from_iter_values(["a", "b", "c"]);
695        let keys = Int8Array::from(vec![Some(2), None, Some(0)]);
696        let dict = DictionaryArray::<Int8Type>::new(keys, Arc::new(values));
697
698        let gc = garbage_collect_dictionary(&dict).unwrap();
699
700        let expected_values = StringArray::from_iter_values(["a", "c"]);
701        let expected_keys = Int8Array::from(vec![Some(1), None, Some(0)]);
702        let expected = DictionaryArray::<Int8Type>::new(expected_keys, Arc::new(expected_values));
703
704        assert_eq!(gc, expected);
705    }
706
707    #[test]
708    fn test_garbage_collect_empty_dictionary() {
709        let values = StringArray::from_iter_values::<&str, _>([]);
710        let keys = Int32Array::from_iter_values([]);
711        let dict = DictionaryArray::<Int32Type>::new(keys, Arc::new(values));
712
713        let gc = garbage_collect_dictionary(&dict).unwrap();
714
715        assert_eq!(gc, dict);
716    }
717
718    #[test]
719    fn test_garbage_collect_dictionary_all_unreferenced() {
720        let values = StringArray::from_iter_values(["a", "b", "c"]);
721        let keys = Int32Array::from(vec![None, None, None]);
722        let dict = DictionaryArray::<Int32Type>::new(keys, Arc::new(values));
723
724        let gc = garbage_collect_dictionary(&dict).unwrap();
725
726        // All keys are null, so dictionary values can be empty
727        let expected_values = StringArray::from_iter_values::<&str, _>([]);
728        let expected_keys = Int32Array::from(vec![None, None, None]);
729        let expected = DictionaryArray::<Int32Type>::new(expected_keys, Arc::new(expected_values));
730
731        assert_eq!(gc, expected);
732    }
733
734    #[test]
735    fn test_merge_strings() {
736        let a = DictionaryArray::<Int32Type>::from_iter(["a", "b", "a", "b", "d", "c", "e"]);
737        let b = DictionaryArray::<Int32Type>::from_iter(["c", "f", "c", "d", "a", "d"]);
738        let merged = merge_dictionary_values(&[&a, &b], None, None).unwrap();
739
740        let values = as_string_array(merged.values.as_ref());
741        let actual: Vec<_> = values.iter().map(Option::unwrap).collect();
742        assert_eq!(&actual, &["a", "b", "d", "c", "e", "f"]);
743
744        assert_eq!(merged.key_mappings.len(), 2);
745        assert_mapping(&merged.key_mappings[0], &[0, 1, 2, 3, 4]);
746        assert_mapping(&merged.key_mappings[1], &[3, 5, 2, 0]);
747
748        let a_slice = a.slice(1, 4);
749        let merged = merge_dictionary_values(&[&a_slice, &b], None, None).unwrap();
750
751        let values = as_string_array(merged.values.as_ref());
752        let actual: Vec<_> = values.iter().map(Option::unwrap).collect();
753        assert_eq!(&actual, &["a", "b", "d", "c", "f"]);
754
755        assert_eq!(merged.key_mappings.len(), 2);
756        assert_mapping(&merged.key_mappings[0], &[0, 1, 2, 0, 0]);
757        assert_mapping(&merged.key_mappings[1], &[3, 4, 2, 0]);
758
759        // Mask out only ["b", "b", "d"] from a
760        let a_mask = BooleanBuffer::from_iter([false, true, false, true, true, false, false]);
761        let b_mask = BooleanBuffer::new_set(b.len());
762        let merged = merge_dictionary_values(&[&a, &b], Some(&[a_mask, b_mask]), None).unwrap();
763
764        let values = as_string_array(merged.values.as_ref());
765        let actual: Vec<_> = values.iter().map(Option::unwrap).collect();
766        assert_eq!(&actual, &["b", "d", "c", "f", "a"]);
767
768        assert_eq!(merged.key_mappings.len(), 2);
769        assert_mapping(&merged.key_mappings[0], &[0, 0, 1, 0, 0]);
770        assert_mapping(&merged.key_mappings[1], &[2, 3, 1, 4]);
771    }
772
773    #[test]
774    fn test_merge_dictionary_values_with_masks() {
775        let first = StringArray::from_iter_values(["a", "unused-a", "shared"]);
776        let second = StringArray::from_iter_values(["shared", "b", "unused-b"]);
777        let values = [&first as &dyn Array, &second as &dyn Array];
778        let masks = [
779            BooleanBuffer::from_iter([true, false, true]),
780            BooleanBuffer::from_iter([true, true, false]),
781        ];
782
783        let merged = merge_dictionary_values_with_masks::<Int32Type>(&values, &masks).unwrap();
784
785        let expected = StringArray::from_iter_values(["a", "shared", "shared", "b"]);
786        assert_eq!(merged.values.as_ref(), &expected);
787        assert_mapping(&merged.key_mappings[0], &[0, 0, 1]);
788        assert_mapping(&merged.key_mappings[1], &[2, 3, 0]);
789    }
790
791    #[test]
792    fn test_merge_nulls() {
793        let buffer = Buffer::from(b"helloworldbingohelloworld");
794        let offsets = OffsetBuffer::from_lengths([5, 5, 5, 5, 5]);
795        let nulls = NullBuffer::from(vec![true, false, true, true, true]);
796        let values = StringArray::new(offsets, buffer, Some(nulls));
797
798        let key_values = vec![1, 2, 3, 1, 8, 2, 3];
799        let key_nulls = NullBuffer::from(vec![true, true, false, true, false, true, true]);
800        let keys = Int32Array::new(key_values.into(), Some(key_nulls));
801        let a = DictionaryArray::new(keys, Arc::new(values));
802        // [NULL, "bingo", NULL, NULL, NULL, "bingo", "hello"]
803
804        let b = DictionaryArray::new(Int32Array::new_null(10), Arc::new(StringArray::new_null(0)));
805
806        let merged = merge_dictionary_values(&[&a, &b], None, None).unwrap();
807        let expected = StringArray::from(vec![None, Some("bingo"), Some("hello")]);
808        assert_eq!(merged.values.as_ref(), &expected);
809        assert_eq!(merged.key_mappings.len(), 2);
810        assert_mapping(&merged.key_mappings[0], &[0, 0, 1, 2, 0]);
811        assert_mapping(&merged.key_mappings[1], &[]);
812    }
813
814    #[test]
815    fn test_merge_keys_smaller() {
816        let values = StringArray::from_iter_values(["a", "b"]);
817        let keys = Int32Array::from_iter_values([1]);
818        let a = DictionaryArray::new(keys, Arc::new(values));
819
820        let merged = merge_dictionary_values(&[&a], None, None).unwrap();
821        let expected = StringArray::from(vec!["b"]);
822        assert_eq!(merged.values.as_ref(), &expected);
823    }
824
825    #[test]
826    fn test_compute_referenced_keys_preserves_dense_order() {
827        let keys = Int32Array::from_iter_values([999, 1, 500, 1, 999]);
828
829        let actual = compute_referenced_keys(keys.values(), None);
830
831        assert_eq!(actual, vec![1, 500, 999]);
832    }
833
834    #[test]
835    fn test_sparse_gate_counts_key_lookups() {
836        let values: ArrayRef = Arc::new(StringArray::from_iter_values(
837            (0..8 * MAX_DENSE_MAPPING_VALUES_PER_ROW).map(|idx| idx.to_string()),
838        ));
839        let is_sparse = |dictionary: &DictionaryArray<Int32Type>,
840                         masks: Option<&[BooleanBuffer]>,
841                         lookups: Option<&[usize]>| {
842            let merged = merge_dictionary_values(&[dictionary], masks, lookups).unwrap();
843            matches!(merged.key_mappings[0], KeyMapping::Sparse(_))
844        };
845
846        // One valid row of ten is over the gate by selected rows, but concat remaps all ten.
847        let keys = Int32Array::from_iter((0..10).map(|idx| (idx == 0).then_some(999)));
848        let mostly_null = DictionaryArray::new(keys, Arc::clone(&values));
849        assert!(!is_sparse(&mostly_null, None, None));
850
851        // Seven lookups stay under the gate, keeping the sparse path.
852        let keys = Int32Array::from_iter_values([999; 7]);
853        let short = DictionaryArray::new(keys, Arc::clone(&values));
854        assert!(is_sparse(&short, None, None));
855
856        // Interleave may look up the selected row more than once.
857        let mask = BooleanBuffer::collect_bool(7, |idx| idx == 0);
858        assert!(!is_sparse(
859            &short,
860            Some(std::slice::from_ref(&mask)),
861            Some(&[10]),
862        ));
863        assert!(is_sparse(&short, Some(&[mask]), Some(&[7])));
864    }
865
866    #[test]
867    fn test_merge_oversized_dictionary_slice() {
868        let values_len = 4 * MAX_DENSE_MAPPING_VALUES_PER_ROW + 1;
869        let values = StringArray::from_iter_values((0..values_len).map(|idx| idx.to_string()));
870        let last = values_len as i32 - 1;
871        let keys = Int32Array::from(vec![Some(last), None, Some(1), Some(last)]);
872        let dictionary = DictionaryArray::new(keys, Arc::new(values));
873
874        let merged = merge_dictionary_values(&[&dictionary], None, None).unwrap();
875        let values = as_string_array(merged.values.as_ref());
876        let actual = values.iter().collect::<Vec<_>>();
877        let last_value = last.to_string();
878
879        assert_eq!(actual, vec![Some("1"), Some(last_value.as_str())]);
880        assert!(matches!(merged.key_mappings[0], KeyMapping::Sparse(_)));
881        assert_eq!(merged.key_mappings[0].get(1), Some(0));
882        assert_eq!(merged.key_mappings[0].get(last as usize), Some(1));
883        assert_eq!(merged.key_mappings[0].get(500), None);
884    }
885
886    #[test]
887    fn test_merge_sparse_dictionary_overflow() {
888        const UNIQUE_VALUES: usize = u8::MAX as usize + 2;
889
890        // One sparse value per dictionary exceeds the 256 values representable by UInt8.
891        let dictionaries = (0..UNIQUE_VALUES)
892            .map(|dictionary_idx| {
893                let values = StringArray::from_iter_values(
894                    (0..=MAX_DENSE_MAPPING_VALUES_PER_ROW)
895                        .map(|value_idx| format!("{dictionary_idx}-{value_idx}")),
896                );
897                let keys = UInt8Array::from_iter_values([0]);
898                DictionaryArray::<UInt8Type>::new(keys, Arc::new(values))
899            })
900            .collect::<Vec<_>>();
901        let dictionaries = dictionaries.iter().collect::<Vec<_>>();
902
903        let result = merge_dictionary_values(&dictionaries, None, None);
904
905        assert!(matches!(
906            result,
907            Err(ArrowError::DictionaryKeyOverflowError)
908        ));
909    }
910}
arrow_select/dictionary.rs

arrow_select/
dictionary.rs
