arrow_select/
dictionary.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Dictionary utilities for Arrow arrays
19
20use std::sync::Arc;
21
22use crate::filter::filter;
23use crate::interleave::interleave;
24use ahash::RandomState;
25use arrow_array::builder::BooleanBufferBuilder;
26use arrow_array::types::{
27    ArrowDictionaryKeyType, ArrowPrimitiveType, BinaryType, ByteArrayType, LargeBinaryType,
28    LargeUtf8Type, Utf8Type,
29};
30use arrow_array::{
31    AnyDictionaryArray, Array, ArrayRef, ArrowNativeTypeOp, BooleanArray, DictionaryArray,
32    GenericByteArray, PrimitiveArray, downcast_dictionary_array,
33};
34use arrow_array::{cast::AsArray, downcast_primitive};
35use arrow_buffer::{ArrowNativeType, BooleanBuffer, ScalarBuffer, ToByteSlice};
36use arrow_schema::{ArrowError, DataType};
37
38/// Garbage collects a [DictionaryArray] by removing unreferenced values.
39///
40/// Returns a new [DictionaryArray] such that there are no values
41/// that are not referenced by at least one key. There may still be duplicate
42/// values.
43///
44/// See also [`garbage_collect_any_dictionary`] if you need to handle multiple dictionary types
45pub fn garbage_collect_dictionary<K: ArrowDictionaryKeyType>(
46    dictionary: &DictionaryArray<K>,
47) -> Result<DictionaryArray<K>, ArrowError> {
48    let keys = dictionary.keys();
49    let values = dictionary.values();
50
51    let mask = dictionary.occupancy();
52
53    // If no work to do, return the original dictionary
54    if mask.count_set_bits() == values.len() {
55        return Ok(dictionary.clone());
56    }
57
58    // Create a mapping from the old keys to the new keys, use a Vec for easy indexing
59    let mut key_remap = vec![K::Native::ZERO; values.len()];
60    for (new_idx, old_idx) in mask.set_indices().enumerate() {
61        key_remap[old_idx] = K::Native::from_usize(new_idx)
62            .expect("new index should fit in K::Native, as old index was in range");
63    }
64
65    // ... and then build the new keys array
66    let new_keys = keys.unary(|key| {
67        key_remap
68            .get(key.as_usize())
69            .copied()
70            // nulls may be present in the keys, and they will have arbitrary value; we don't care
71            // and can safely return zero
72            .unwrap_or(K::Native::ZERO)
73    });
74
75    // Create a new values array by filtering using the mask
76    let values = filter(dictionary.values(), &BooleanArray::new(mask, None))?;
77
78    DictionaryArray::try_new(new_keys, values)
79}
80
81/// Equivalent to [`garbage_collect_dictionary`] but without requiring casting to a specific key type.
82pub fn garbage_collect_any_dictionary(
83    dictionary: &dyn AnyDictionaryArray,
84) -> Result<ArrayRef, ArrowError> {
85    let dictionary: &dyn Array = dictionary;
86    downcast_dictionary_array!(
87        dictionary => garbage_collect_dictionary(dictionary).map(|dict| Arc::new(dict) as ArrayRef),
88        _ => unreachable!("have a dictionary array")
89    )
90}
91
92/// A best effort interner that maintains a fixed number of buckets
93/// and interns keys based on their hash value
94///
95/// Hash collisions will result in replacement
96struct Interner<'a, V> {
97    state: RandomState,
98    buckets: Vec<Option<InternerBucket<'a, V>>>,
99    shift: u32,
100}
101
102/// A single bucket in [`Interner`].
103type InternerBucket<'a, V> = (Option<&'a [u8]>, V);
104
105impl<'a, V> Interner<'a, V> {
106    /// Capacity controls the number of unique buckets allocated within the Interner
107    ///
108    /// A larger capacity reduces the probability of hash collisions, and should be set
109    /// based on an approximation of the upper bound of unique values
110    fn new(capacity: usize) -> Self {
111        // Add additional buckets to help reduce collisions
112        let shift = (capacity as u64 + 128).leading_zeros();
113        let num_buckets = (u64::MAX >> shift) as usize;
114        let buckets = (0..num_buckets.saturating_add(1)).map(|_| None).collect();
115        Self {
116            // A fixed seed to ensure deterministic behaviour
117            state: RandomState::with_seeds(0, 0, 0, 0),
118            buckets,
119            shift,
120        }
121    }
122
123    fn intern<F: FnOnce() -> Result<V, E>, E>(
124        &mut self,
125        new: Option<&'a [u8]>,
126        f: F,
127    ) -> Result<&V, E> {
128        let hash = self.state.hash_one(new);
129        let bucket_idx = hash >> self.shift;
130        Ok(match &mut self.buckets[bucket_idx as usize] {
131            Some((current, v)) => {
132                if *current != new {
133                    *v = f()?;
134                    *current = new;
135                }
136                v
137            }
138            slot => &slot.insert((new, f()?)).1,
139        })
140    }
141}
142
143pub(crate) struct MergedDictionaries<K: ArrowDictionaryKeyType> {
144    /// Provides `key_mappings[`array_idx`][`old_key`] -> new_key`
145    pub key_mappings: Vec<Vec<K::Native>>,
146    /// The new values
147    pub values: ArrayRef,
148}
149
150/// Performs a cheap, pointer-based comparison of two byte array
151///
152/// See [`ScalarBuffer::ptr_eq`]
153fn bytes_ptr_eq<T: ByteArrayType>(a: &dyn Array, b: &dyn Array) -> bool {
154    match (a.as_bytes_opt::<T>(), b.as_bytes_opt::<T>()) {
155        (Some(a), Some(b)) => {
156            let values_eq = a.values().ptr_eq(b.values()) && a.offsets().ptr_eq(b.offsets());
157            match (a.nulls(), b.nulls()) {
158                (Some(a), Some(b)) => values_eq && a.inner().ptr_eq(b.inner()),
159                (None, None) => values_eq,
160                _ => false,
161            }
162        }
163        _ => false,
164    }
165}
166
167/// A type-erased function that compares two array for pointer equality
168type PtrEq = fn(&dyn Array, &dyn Array) -> bool;
169
170/// A weak heuristic of whether to merge dictionary values that aims to only
171/// perform the expensive merge computation when it is likely to yield at least
172/// some return over the naive approach used by MutableArrayData
173///
174/// `len` is the total length of the merged output
175///
176/// Returns `(should_merge, has_overflow)` where:
177/// - `should_merge`: whether dictionary values should be merged
178/// - `has_overflow`: whether the combined dictionary values would overflow the key type
179pub(crate) fn should_merge_dictionary_values<K: ArrowDictionaryKeyType>(
180    dictionaries: &[&DictionaryArray<K>],
181    len: usize,
182) -> (bool, bool) {
183    use DataType::*;
184    let first_values = dictionaries[0].values().as_ref();
185    let ptr_eq: PtrEq = match first_values.data_type() {
186        Utf8 => bytes_ptr_eq::<Utf8Type>,
187        LargeUtf8 => bytes_ptr_eq::<LargeUtf8Type>,
188        Binary => bytes_ptr_eq::<BinaryType>,
189        LargeBinary => bytes_ptr_eq::<LargeBinaryType>,
190        dt => {
191            if !dt.is_primitive() {
192                return (
193                    false,
194                    K::Native::from_usize(dictionaries.iter().map(|d| d.values().len()).sum())
195                        .is_none(),
196                );
197            }
198            |a, b| a.to_data().ptr_eq(&b.to_data())
199        }
200    };
201
202    let mut single_dictionary = true;
203    let mut total_values = first_values.len();
204    for dict in dictionaries.iter().skip(1) {
205        let values = dict.values().as_ref();
206        total_values += values.len();
207        if single_dictionary {
208            single_dictionary = ptr_eq(first_values, values)
209        }
210    }
211
212    let overflow = K::Native::from_usize(total_values).is_none();
213    let values_exceed_length = total_values >= len;
214
215    (
216        !single_dictionary && (overflow || values_exceed_length),
217        overflow,
218    )
219}
220
221/// Given an array of dictionaries and an optional key mask compute a values array
222/// containing referenced values, along with mappings from the [`DictionaryArray`]
223/// keys to the new keys within this values array. Best-effort will be made to ensure
224/// that the dictionary values are unique
225///
226/// This method is meant to be very fast and the output dictionary values
227/// may not be unique, unlike `GenericByteDictionaryBuilder` which is slower
228/// but produces unique values
229pub(crate) fn merge_dictionary_values<K: ArrowDictionaryKeyType>(
230    dictionaries: &[&DictionaryArray<K>],
231    masks: Option<&[BooleanBuffer]>,
232) -> Result<MergedDictionaries<K>, ArrowError> {
233    let mut num_values = 0;
234
235    let mut values_arrays = Vec::with_capacity(dictionaries.len());
236    let mut value_slices = Vec::with_capacity(dictionaries.len());
237
238    for (idx, dictionary) in dictionaries.iter().enumerate() {
239        let mask = masks.and_then(|m| m.get(idx));
240        let key_mask_owned;
241        let key_mask = match (dictionary.nulls(), mask) {
242            (Some(n), None) => Some(n.inner()),
243            (None, Some(n)) => Some(n),
244            (Some(n), Some(m)) => {
245                key_mask_owned = n.inner() & m;
246                Some(&key_mask_owned)
247            }
248            (None, None) => None,
249        };
250        let keys = dictionary.keys().values();
251        let values = dictionary.values().as_ref();
252        let values_mask = compute_values_mask(keys, key_mask, values.len());
253
254        let masked_values = get_masked_values(values, &values_mask);
255        num_values += masked_values.len();
256        value_slices.push(masked_values);
257        values_arrays.push(values)
258    }
259
260    // Map from value to new index
261    let mut interner = Interner::new(num_values);
262    // Interleave indices for new values array
263    let mut indices = Vec::with_capacity(num_values);
264
265    // Compute the mapping for each dictionary
266    let key_mappings = dictionaries
267        .iter()
268        .enumerate()
269        .zip(value_slices)
270        .map(|((dictionary_idx, dictionary), values)| {
271            let zero = K::Native::from_usize(0).unwrap();
272            let mut mapping = vec![zero; dictionary.values().len()];
273
274            for (value_idx, value) in values {
275                mapping[value_idx] =
276                    *interner.intern(value, || match K::Native::from_usize(indices.len()) {
277                        Some(idx) => {
278                            indices.push((dictionary_idx, value_idx));
279                            Ok(idx)
280                        }
281                        None => Err(ArrowError::DictionaryKeyOverflowError),
282                    })?;
283            }
284            Ok(mapping)
285        })
286        .collect::<Result<Vec<_>, ArrowError>>()?;
287
288    Ok(MergedDictionaries {
289        key_mappings,
290        values: interleave(&values_arrays, &indices)?,
291    })
292}
293
294/// Return a mask identifying the values that are referenced by keys in `dictionary`
295/// at the positions indicated by `selection`
296fn compute_values_mask<K: ArrowNativeType>(
297    keys: &ScalarBuffer<K>,
298    mask: Option<&BooleanBuffer>,
299    max_key: usize,
300) -> BooleanBuffer {
301    let mut builder = BooleanBufferBuilder::new(max_key);
302    builder.advance(max_key);
303
304    match mask {
305        Some(n) => n
306            .set_indices()
307            .for_each(|idx| builder.set_bit(keys[idx].as_usize(), true)),
308        None => keys
309            .iter()
310            .for_each(|k| builder.set_bit(k.as_usize(), true)),
311    }
312    builder.finish()
313}
314
315/// Process primitive array values to bytes
316fn masked_primitives_to_bytes<'a, T: ArrowPrimitiveType>(
317    array: &'a PrimitiveArray<T>,
318    mask: &BooleanBuffer,
319) -> Vec<(usize, Option<&'a [u8]>)>
320where
321    T::Native: ToByteSlice,
322{
323    let mut out = Vec::with_capacity(mask.count_set_bits());
324    let values = array.values();
325    for idx in mask.set_indices() {
326        out.push((
327            idx,
328            array.is_valid(idx).then_some(values[idx].to_byte_slice()),
329        ))
330    }
331    out
332}
333
334macro_rules! masked_primitive_to_bytes_helper {
335    ($t:ty, $array:expr, $mask:expr) => {
336        masked_primitives_to_bytes::<$t>($array.as_primitive(), $mask)
337    };
338}
339
340/// Return a Vec containing for each set index in `mask`, the index and byte value of that index
341fn get_masked_values<'a>(
342    array: &'a dyn Array,
343    mask: &BooleanBuffer,
344) -> Vec<(usize, Option<&'a [u8]>)> {
345    downcast_primitive! {
346        array.data_type() => (masked_primitive_to_bytes_helper, array, mask),
347        DataType::Utf8 => masked_bytes(array.as_string::<i32>(), mask),
348        DataType::LargeUtf8 => masked_bytes(array.as_string::<i64>(), mask),
349        DataType::Binary => masked_bytes(array.as_binary::<i32>(), mask),
350        DataType::LargeBinary => masked_bytes(array.as_binary::<i64>(), mask),
351        _ => unimplemented!("Dictionary merging for type {} is not implemented", array.data_type()),
352    }
353}
354
355/// Compute [`get_masked_values`] for a [`GenericByteArray`]
356///
357/// Note: this does not check the null mask and will return values contained in null slots
358fn masked_bytes<'a, T: ByteArrayType>(
359    array: &'a GenericByteArray<T>,
360    mask: &BooleanBuffer,
361) -> Vec<(usize, Option<&'a [u8]>)> {
362    let mut out = Vec::with_capacity(mask.count_set_bits());
363    for idx in mask.set_indices() {
364        out.push((
365            idx,
366            array.is_valid(idx).then_some(array.value(idx).as_ref()),
367        ))
368    }
369    out
370}
371
372#[cfg(test)]
373mod tests {
374    use super::*;
375
376    use arrow_array::cast::as_string_array;
377    use arrow_array::types::Int8Type;
378    use arrow_array::types::Int32Type;
379    use arrow_array::{DictionaryArray, Int8Array, Int32Array, StringArray};
380    use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer, OffsetBuffer};
381    use std::sync::Arc;
382
383    #[test]
384    fn test_garbage_collect_i32_dictionary() {
385        let values = StringArray::from_iter_values(["a", "b", "c", "d"]);
386        let keys = Int32Array::from_iter_values([0, 1, 1, 3, 0, 0, 1]);
387        let dict = DictionaryArray::<Int32Type>::new(keys, Arc::new(values));
388
389        // Only "a", "b", "d" are referenced, "c" is not
390        let gc = garbage_collect_dictionary(&dict).unwrap();
391
392        let expected_values = StringArray::from_iter_values(["a", "b", "d"]);
393        let expected_keys = Int32Array::from_iter_values([0, 1, 1, 2, 0, 0, 1]);
394        let expected = DictionaryArray::<Int32Type>::new(expected_keys, Arc::new(expected_values));
395
396        assert_eq!(gc, expected);
397    }
398
399    #[test]
400    fn test_garbage_collect_any_dictionary() {
401        let values = StringArray::from_iter_values(["a", "b", "c", "d"]);
402        let keys = Int32Array::from_iter_values([0, 1, 1, 3, 0, 0, 1]);
403        let dict = DictionaryArray::<Int32Type>::new(keys, Arc::new(values));
404
405        let gc = garbage_collect_any_dictionary(&dict).unwrap();
406
407        let expected_values = StringArray::from_iter_values(["a", "b", "d"]);
408        let expected_keys = Int32Array::from_iter_values([0, 1, 1, 2, 0, 0, 1]);
409        let expected = DictionaryArray::<Int32Type>::new(expected_keys, Arc::new(expected_values));
410
411        assert_eq!(gc.as_ref(), &expected);
412    }
413
414    #[test]
415    fn test_garbage_collect_with_nulls() {
416        let values = StringArray::from_iter_values(["a", "b", "c"]);
417        let keys = Int8Array::from(vec![Some(2), None, Some(0)]);
418        let dict = DictionaryArray::<Int8Type>::new(keys, Arc::new(values));
419
420        let gc = garbage_collect_dictionary(&dict).unwrap();
421
422        let expected_values = StringArray::from_iter_values(["a", "c"]);
423        let expected_keys = Int8Array::from(vec![Some(1), None, Some(0)]);
424        let expected = DictionaryArray::<Int8Type>::new(expected_keys, Arc::new(expected_values));
425
426        assert_eq!(gc, expected);
427    }
428
429    #[test]
430    fn test_garbage_collect_empty_dictionary() {
431        let values = StringArray::from_iter_values::<&str, _>([]);
432        let keys = Int32Array::from_iter_values([]);
433        let dict = DictionaryArray::<Int32Type>::new(keys, Arc::new(values));
434
435        let gc = garbage_collect_dictionary(&dict).unwrap();
436
437        assert_eq!(gc, dict);
438    }
439
440    #[test]
441    fn test_garbage_collect_dictionary_all_unreferenced() {
442        let values = StringArray::from_iter_values(["a", "b", "c"]);
443        let keys = Int32Array::from(vec![None, None, None]);
444        let dict = DictionaryArray::<Int32Type>::new(keys, Arc::new(values));
445
446        let gc = garbage_collect_dictionary(&dict).unwrap();
447
448        // All keys are null, so dictionary values can be empty
449        let expected_values = StringArray::from_iter_values::<&str, _>([]);
450        let expected_keys = Int32Array::from(vec![None, None, None]);
451        let expected = DictionaryArray::<Int32Type>::new(expected_keys, Arc::new(expected_values));
452
453        assert_eq!(gc, expected);
454    }
455
456    #[test]
457    fn test_merge_strings() {
458        let a = DictionaryArray::<Int32Type>::from_iter(["a", "b", "a", "b", "d", "c", "e"]);
459        let b = DictionaryArray::<Int32Type>::from_iter(["c", "f", "c", "d", "a", "d"]);
460        let merged = merge_dictionary_values(&[&a, &b], None).unwrap();
461
462        let values = as_string_array(merged.values.as_ref());
463        let actual: Vec<_> = values.iter().map(Option::unwrap).collect();
464        assert_eq!(&actual, &["a", "b", "d", "c", "e", "f"]);
465
466        assert_eq!(merged.key_mappings.len(), 2);
467        assert_eq!(&merged.key_mappings[0], &[0, 1, 2, 3, 4]);
468        assert_eq!(&merged.key_mappings[1], &[3, 5, 2, 0]);
469
470        let a_slice = a.slice(1, 4);
471        let merged = merge_dictionary_values(&[&a_slice, &b], None).unwrap();
472
473        let values = as_string_array(merged.values.as_ref());
474        let actual: Vec<_> = values.iter().map(Option::unwrap).collect();
475        assert_eq!(&actual, &["a", "b", "d", "c", "f"]);
476
477        assert_eq!(merged.key_mappings.len(), 2);
478        assert_eq!(&merged.key_mappings[0], &[0, 1, 2, 0, 0]);
479        assert_eq!(&merged.key_mappings[1], &[3, 4, 2, 0]);
480
481        // Mask out only ["b", "b", "d"] from a
482        let a_mask = BooleanBuffer::from_iter([false, true, false, true, true, false, false]);
483        let b_mask = BooleanBuffer::new_set(b.len());
484        let merged = merge_dictionary_values(&[&a, &b], Some(&[a_mask, b_mask])).unwrap();
485
486        let values = as_string_array(merged.values.as_ref());
487        let actual: Vec<_> = values.iter().map(Option::unwrap).collect();
488        assert_eq!(&actual, &["b", "d", "c", "f", "a"]);
489
490        assert_eq!(merged.key_mappings.len(), 2);
491        assert_eq!(&merged.key_mappings[0], &[0, 0, 1, 0, 0]);
492        assert_eq!(&merged.key_mappings[1], &[2, 3, 1, 4]);
493    }
494
495    #[test]
496    fn test_merge_nulls() {
497        let buffer = Buffer::from(b"helloworldbingohelloworld");
498        let offsets = OffsetBuffer::from_lengths([5, 5, 5, 5, 5]);
499        let nulls = NullBuffer::from(vec![true, false, true, true, true]);
500        let values = StringArray::new(offsets, buffer, Some(nulls));
501
502        let key_values = vec![1, 2, 3, 1, 8, 2, 3];
503        let key_nulls = NullBuffer::from(vec![true, true, false, true, false, true, true]);
504        let keys = Int32Array::new(key_values.into(), Some(key_nulls));
505        let a = DictionaryArray::new(keys, Arc::new(values));
506        // [NULL, "bingo", NULL, NULL, NULL, "bingo", "hello"]
507
508        let b = DictionaryArray::new(Int32Array::new_null(10), Arc::new(StringArray::new_null(0)));
509
510        let merged = merge_dictionary_values(&[&a, &b], None).unwrap();
511        let expected = StringArray::from(vec![None, Some("bingo"), Some("hello")]);
512        assert_eq!(merged.values.as_ref(), &expected);
513        assert_eq!(merged.key_mappings.len(), 2);
514        assert_eq!(&merged.key_mappings[0], &[0, 0, 1, 2, 0]);
515        assert_eq!(&merged.key_mappings[1], &[] as &[i32; 0]);
516    }
517
518    #[test]
519    fn test_merge_keys_smaller() {
520        let values = StringArray::from_iter_values(["a", "b"]);
521        let keys = Int32Array::from_iter_values([1]);
522        let a = DictionaryArray::new(keys, Arc::new(values));
523
524        let merged = merge_dictionary_values(&[&a], None).unwrap();
525        let expected = StringArray::from(vec!["b"]);
526        assert_eq!(merged.values.as_ref(), &expected);
527    }
528}
arrow_select/dictionary.rs

arrow_select/
dictionary.rs
