arrow_array/builder/
generic_bytes_dictionary_builder.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
5use crate::builder::{ArrayBuilder, GenericByteBuilder, PrimitiveBuilder};
6use crate::types::{ArrowDictionaryKeyType, ByteArrayType, GenericBinaryType, GenericStringType};
7use crate::{
8    Array, ArrayRef, DictionaryArray, GenericByteArray, PrimitiveArray, TypedDictionaryArray,
9};
10use arrow_buffer::ArrowNativeType;
11use arrow_schema::{ArrowError, DataType};
12use hashbrown::HashTable;
13use num_traits::NumCast;
14use std::any::Any;
15use std::sync::Arc;
16
17/// Builder for [`DictionaryArray`] of [`GenericByteArray`]
18///
19/// For example to map a set of byte indices to String values. Note that
20/// the use of a `HashMap` here will not scale to very large arrays or
21/// result in an ordered dictionary.
22#[derive(Debug)]
23pub struct GenericByteDictionaryBuilder<K, T>
24where
25    K: ArrowDictionaryKeyType,
26    T: ByteArrayType,
27{
28    state: ahash::RandomState,
29    dedup: HashTable<usize>,
30
31    keys_builder: PrimitiveBuilder<K>,
32    values_builder: GenericByteBuilder<T>,
33}
34
35impl<K, T> Default for GenericByteDictionaryBuilder<K, T>
36where
37    K: ArrowDictionaryKeyType,
38    T: ByteArrayType,
39{
40    fn default() -> Self {
41        Self::new()
42    }
43}
44
45impl<K, T> GenericByteDictionaryBuilder<K, T>
46where
47    K: ArrowDictionaryKeyType,
48    T: ByteArrayType,
49{
50    /// Creates a new `GenericByteDictionaryBuilder`
51    pub fn new() -> Self {
52        let keys_builder = PrimitiveBuilder::new();
53        let values_builder = GenericByteBuilder::<T>::new();
54        Self {
55            state: Default::default(),
56            dedup: HashTable::with_capacity(keys_builder.capacity()),
57            keys_builder,
58            values_builder,
59        }
60    }
61
62    /// Creates a new `GenericByteDictionaryBuilder` with the provided capacities
63    ///
64    /// `keys_capacity`: the number of keys, i.e. length of array to build
65    /// `value_capacity`: the number of distinct dictionary values, i.e. size of dictionary
66    /// `data_capacity`: the total number of bytes of all distinct bytes in the dictionary
67    pub fn with_capacity(
68        keys_capacity: usize,
69        value_capacity: usize,
70        data_capacity: usize,
71    ) -> Self {
72        Self {
73            state: Default::default(),
74            dedup: HashTable::with_capacity(value_capacity),
75            keys_builder: PrimitiveBuilder::with_capacity(keys_capacity),
76            values_builder: GenericByteBuilder::<T>::with_capacity(value_capacity, data_capacity),
77        }
78    }
79
80    /// Creates a new `GenericByteDictionaryBuilder` from a keys capacity and a dictionary
81    /// which is initialized with the given values.
82    /// The indices of those dictionary values are used as keys.
83    ///
84    /// # Example
85    ///
86    /// ```
87    /// # use arrow_array::builder::StringDictionaryBuilder;
88    /// # use arrow_array::{Int16Array, StringArray};
89    ///
90    /// let dictionary_values = StringArray::from(vec![None, Some("abc"), Some("def")]);
91    ///
92    /// let mut builder = StringDictionaryBuilder::new_with_dictionary(3, &dictionary_values).unwrap();
93    /// builder.append("def").unwrap();
94    /// builder.append_null();
95    /// builder.append("abc").unwrap();
96    ///
97    /// let dictionary_array = builder.finish();
98    ///
99    /// let keys = dictionary_array.keys();
100    ///
101    /// assert_eq!(keys, &Int16Array::from(vec![Some(2), None, Some(1)]));
102    /// ```
103    pub fn new_with_dictionary(
104        keys_capacity: usize,
105        dictionary_values: &GenericByteArray<T>,
106    ) -> Result<Self, ArrowError> {
107        let state = ahash::RandomState::default();
108        let dict_len = dictionary_values.len();
109
110        let mut dedup = HashTable::with_capacity(dict_len);
111
112        let values_len = dictionary_values.value_data().len();
113        let mut values_builder = GenericByteBuilder::<T>::with_capacity(dict_len, values_len);
114
115        K::Native::from_usize(dictionary_values.len())
116            .ok_or(ArrowError::DictionaryKeyOverflowError)?;
117
118        for (idx, maybe_value) in dictionary_values.iter().enumerate() {
119            match maybe_value {
120                Some(value) => {
121                    let value_bytes: &[u8] = value.as_ref();
122                    let hash = state.hash_one(value_bytes);
123
124                    dedup
125                        .entry(
126                            hash,
127                            |idx: &usize| value_bytes == get_bytes(&values_builder, *idx),
128                            |idx: &usize| state.hash_one(get_bytes(&values_builder, *idx)),
129                        )
130                        .or_insert(idx);
131
132                    values_builder.append_value(value);
133                }
134                None => values_builder.append_null(),
135            }
136        }
137
138        Ok(Self {
139            state,
140            dedup,
141            keys_builder: PrimitiveBuilder::with_capacity(keys_capacity),
142            values_builder,
143        })
144    }
145
146    /// Creates a new `GenericByteDictionaryBuilder` from the existing builder with the same
147    /// keys and values, but with a new data type for the keys.
148    ///
149    /// # Example
150    /// ```
151    /// #
152    /// # use arrow_array::builder::StringDictionaryBuilder;
153    /// # use arrow_array::types::{UInt8Type, UInt16Type};
154    /// # use arrow_array::UInt16Array;
155    /// # use arrow_schema::ArrowError;
156    ///
157    /// let mut u8_keyed_builder = StringDictionaryBuilder::<UInt8Type>::new();
158    ///
159    /// // appending too many values causes the dictionary to overflow
160    /// for i in 0..256 {
161    ///     u8_keyed_builder.append_value(format!("{}", i));
162    /// }
163    /// let result = u8_keyed_builder.append("256");
164    /// assert!(matches!(result, Err(ArrowError::DictionaryKeyOverflowError{})));
165    ///
166    /// // we need to upgrade to a larger key type
167    /// let mut u16_keyed_builder = StringDictionaryBuilder::<UInt16Type>::try_new_from_builder(u8_keyed_builder).unwrap();
168    /// let dictionary_array = u16_keyed_builder.finish();
169    /// let keys = dictionary_array.keys();
170    ///
171    /// assert_eq!(keys, &UInt16Array::from_iter(0..256));
172    /// ```
173    pub fn try_new_from_builder<K2>(
174        mut source: GenericByteDictionaryBuilder<K2, T>,
175    ) -> Result<Self, ArrowError>
176    where
177        K::Native: NumCast,
178        K2: ArrowDictionaryKeyType,
179        K2::Native: NumCast,
180    {
181        let state = source.state;
182        let dedup = source.dedup;
183        let values_builder = source.values_builder;
184
185        let source_keys = source.keys_builder.finish();
186        let new_keys: PrimitiveArray<K> = source_keys.try_unary(|value| {
187            num_traits::cast::cast::<K2::Native, K::Native>(value).ok_or_else(|| {
188                ArrowError::CastError(format!(
189                    "Can't cast dictionary keys from source type {:?} to type {:?}",
190                    K2::DATA_TYPE,
191                    K::DATA_TYPE
192                ))
193            })
194        })?;
195
196        // drop source key here because currently source_keys and new_keys are holding reference to
197        // the same underlying null_buffer. Below we want to call new_keys.into_builder() it must
198        // be the only reference holder.
199        drop(source_keys);
200
201        Ok(Self {
202            state,
203            dedup,
204            keys_builder: new_keys.into_builder().map_err(|_| {
205                ArrowError::ComputeError(
206                    "Internal Error: the keys just derived from the source builder are \
207                     unexpectedly shared, so they cannot be reused as a builder"
208                        .to_string(),
209                )
210            })?,
211            values_builder,
212        })
213    }
214}
215
216impl<K, T> ArrayBuilder for GenericByteDictionaryBuilder<K, T>
217where
218    K: ArrowDictionaryKeyType,
219    T: ByteArrayType,
220{
221    /// Returns the builder as an non-mutable `Any` reference.
222    fn as_any(&self) -> &dyn Any {
223        self
224    }
225
226    /// Returns the builder as an mutable `Any` reference.
227    fn as_any_mut(&mut self) -> &mut dyn Any {
228        self
229    }
230
231    /// Returns the boxed builder as a box of `Any`.
232    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
233        self
234    }
235
236    /// Returns the number of array slots in the builder
237    fn len(&self) -> usize {
238        self.keys_builder.len()
239    }
240
241    /// Builds the array and reset this builder.
242    fn finish(&mut self) -> ArrayRef {
243        Arc::new(self.finish())
244    }
245
246    /// Builds the array without resetting the builder.
247    fn finish_cloned(&self) -> ArrayRef {
248        Arc::new(self.finish_cloned())
249    }
250
251    fn finish_preserve_values(&mut self) -> ArrayRef {
252        Arc::new(self.finish_preserve_values())
253    }
254}
255
256impl<K, T> GenericByteDictionaryBuilder<K, T>
257where
258    K: ArrowDictionaryKeyType,
259    T: ByteArrayType,
260{
261    fn get_or_insert_key(&mut self, value: impl AsRef<T::Native>) -> Result<K::Native, ArrowError> {
262        let value_native: &T::Native = value.as_ref();
263        let value_bytes: &[u8] = value_native.as_ref();
264
265        let state = &self.state;
266        let storage = &mut self.values_builder;
267        let hash = state.hash_one(value_bytes);
268
269        let idx = *self
270            .dedup
271            .entry(
272                hash,
273                |idx| value_bytes == get_bytes(storage, *idx),
274                |idx| state.hash_one(get_bytes(storage, *idx)),
275            )
276            .or_insert_with(|| {
277                let idx = storage.len();
278                storage.append_value(value);
279                idx
280            })
281            .get();
282
283        let key = K::Native::from_usize(idx).ok_or(ArrowError::DictionaryKeyOverflowError)?;
284
285        Ok(key)
286    }
287
288    /// Append a value to the array. Return an existing index
289    /// if already present in the values array or a new index if the
290    /// value is appended to the values array.
291    ///
292    /// Returns an error if the new index would overflow the key type.
293    pub fn append(&mut self, value: impl AsRef<T::Native>) -> Result<K::Native, ArrowError> {
294        let key = self.get_or_insert_key(value)?;
295        self.keys_builder.append_value(key);
296        Ok(key)
297    }
298
299    /// Append a value multiple times to the array.
300    /// This is the same as `append` but allows to append the same value multiple times without doing multiple lookups.
301    ///
302    /// Returns an error if the new index would overflow the key type.
303    pub fn append_n(
304        &mut self,
305        value: impl AsRef<T::Native>,
306        count: usize,
307    ) -> Result<K::Native, ArrowError> {
308        let key = self.get_or_insert_key(value)?;
309        self.keys_builder.append_value_n(key, count);
310        Ok(key)
311    }
312
313    /// Infallibly append a value to this builder
314    ///
315    /// # Panics
316    ///
317    /// Panics if the resulting length of the dictionary values array would exceed `T::Native::MAX`
318    pub fn append_value(&mut self, value: impl AsRef<T::Native>) {
319        self.append(value).expect("dictionary key overflow");
320    }
321
322    /// Infallibly append a value to this builder repeatedly `count` times.
323    /// This is the same as `append_value` but allows to append the same value multiple times without doing multiple lookups.
324    ///
325    /// # Panics
326    ///
327    /// Panics if the resulting length of the dictionary values array would exceed `T::Native::MAX`
328    pub fn append_values(&mut self, value: impl AsRef<T::Native>, count: usize) {
329        self.append_n(value, count)
330            .expect("dictionary key overflow");
331    }
332
333    /// Appends a null slot into the builder
334    #[inline]
335    pub fn append_null(&mut self) {
336        self.keys_builder.append_null()
337    }
338
339    /// Infallibly append `n` null slots into the builder
340    #[inline]
341    pub fn append_nulls(&mut self, n: usize) {
342        self.keys_builder.append_nulls(n)
343    }
344
345    /// Append an `Option` value into the builder
346    ///
347    /// # Panics
348    ///
349    /// Panics if the resulting length of the dictionary values array would exceed `T::Native::MAX`
350    #[inline]
351    pub fn append_option(&mut self, value: Option<impl AsRef<T::Native>>) {
352        match value {
353            None => self.append_null(),
354            Some(v) => self.append_value(v),
355        }
356    }
357
358    /// Append an `Option` value into the builder repeatedly `count` times.
359    /// This is the same as `append_option` but allows to append the same value multiple times without doing multiple lookups.
360    ///
361    /// # Panics
362    ///
363    /// Panics if the resulting length of the dictionary values array would exceed `T::Native::MAX`
364    pub fn append_options(&mut self, value: Option<impl AsRef<T::Native>>, count: usize) {
365        match value {
366            None => self.keys_builder.append_nulls(count),
367            Some(v) => self.append_values(v, count),
368        }
369    }
370
371    /// Append all values from a [`GenericByteArray`] into this builder.
372    ///
373    /// This is more efficient than calling [`Self::append`] in a loop because
374    /// it accesses the raw offset/value buffers directly and bulk-writes the
375    /// resolved keys with a single `append_slice` / `append_values` call.
376    ///
377    /// Returns an error if any new dictionary index would overflow the key type.
378    pub fn append_array(&mut self, array: &GenericByteArray<T>) -> Result<(), ArrowError> {
379        let row_count = array.len();
380        if row_count == 0 {
381            return Ok(());
382        }
383        let offsets = array.value_offsets();
384        let raw_data = array.value_data();
385
386        match array.nulls() {
387            None => {
388                let mut key_buf: Vec<K::Native> = Vec::with_capacity(row_count);
389                for row_idx in 0..row_count {
390                    key_buf.push(resolve_key::<K, T>(
391                        row_idx,
392                        offsets,
393                        raw_data,
394                        array,
395                        &mut self.dedup,
396                        &self.state,
397                        &mut self.values_builder,
398                    )?);
399                }
400                self.keys_builder.append_slice(&key_buf);
401            }
402            Some(nulls) => {
403                let mut key_buf: Vec<K::Native> = Vec::with_capacity(row_count);
404                let mut valid_buf: Vec<bool> = Vec::with_capacity(row_count);
405                for row_idx in 0..row_count {
406                    if nulls.is_null(row_idx) {
407                        key_buf.push(K::Native::usize_as(0));
408                        valid_buf.push(false);
409                    } else {
410                        key_buf.push(resolve_key::<K, T>(
411                            row_idx,
412                            offsets,
413                            raw_data,
414                            array,
415                            &mut self.dedup,
416                            &self.state,
417                            &mut self.values_builder,
418                        )?);
419                        valid_buf.push(true);
420                    }
421                }
422                self.keys_builder.append_values(&key_buf, &valid_buf);
423            }
424        }
425        Ok(())
426    }
427
428    /// Extends builder with an existing dictionary array.
429    ///
430    /// This is the same as [`Self::extend`] but is faster as it translates
431    /// the dictionary values once rather than doing a lookup for each item in the iterator
432    ///
433    /// when dictionary values are null (the actual mapped values) the keys are null
434    ///
435    pub fn extend_dictionary(
436        &mut self,
437        dictionary: &TypedDictionaryArray<K, GenericByteArray<T>>,
438    ) -> Result<(), ArrowError> {
439        let values = dictionary.values();
440
441        let v_len = values.len();
442        let k_len = dictionary.keys().len();
443        if v_len == 0 && k_len == 0 {
444            return Ok(());
445        }
446
447        // All nulls
448        if v_len == 0 {
449            self.append_nulls(k_len);
450            return Ok(());
451        }
452
453        if k_len == 0 {
454            return Err(ArrowError::InvalidArgumentError(
455                "Dictionary keys should not be empty when values are not empty".to_string(),
456            ));
457        }
458
459        // Orphan values will be carried over to the new dictionary
460        let mapped_values = values
461            .iter()
462            // Dictionary values can technically be null, so we need to handle that
463            .map(|dict_value| {
464                dict_value
465                    .map(|dict_value| self.get_or_insert_key(dict_value))
466                    .transpose()
467            })
468            .collect::<Result<Vec<_>, _>>()?;
469
470        // Just insert the keys without additional lookups
471        dictionary.keys().iter().for_each(|key| match key {
472            None => self.append_null(),
473            Some(original_dict_index) => {
474                let index = original_dict_index.as_usize().min(v_len - 1);
475                match mapped_values[index] {
476                    None => self.append_null(),
477                    Some(mapped_value) => self.keys_builder.append_value(mapped_value),
478                }
479            }
480        });
481
482        Ok(())
483    }
484
485    /// Builds the `DictionaryArray` and reset this builder.
486    pub fn finish(&mut self) -> DictionaryArray<K> {
487        self.dedup.clear();
488        let values = self.values_builder.finish();
489        let keys = self.keys_builder.finish();
490
491        let data_type = DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(T::DATA_TYPE));
492
493        let builder = keys
494            .into_data()
495            .into_builder()
496            .data_type(data_type)
497            .child_data(vec![values.into_data()]);
498
499        // SAFETY: builder is constructed from valid key/value arrays produced by the builder
500        DictionaryArray::from(unsafe { builder.build_unchecked() })
501    }
502
503    /// Builds the `DictionaryArray` without resetting the builder.
504    pub fn finish_cloned(&self) -> DictionaryArray<K> {
505        let values = self.values_builder.finish_cloned();
506        let keys = self.keys_builder.finish_cloned();
507
508        let data_type = DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(T::DATA_TYPE));
509
510        let builder = keys
511            .into_data()
512            .into_builder()
513            .data_type(data_type)
514            .child_data(vec![values.into_data()]);
515
516        // SAFETY: builder is constructed from valid key/value arrays produced by the builder
517        DictionaryArray::from(unsafe { builder.build_unchecked() })
518    }
519
520    /// Builds the `DictionaryArray` without resetting the values builder or
521    /// the internal de-duplication map.
522    ///
523    /// The advantage of doing this is that the values will represent the entire
524    /// set of what has been built so-far by this builder and ensures
525    /// consistency in the assignment of keys to values across multiple calls
526    /// to `finish_preserve_values`. This enables ipc writers to efficiently
527    /// emit delta dictionaries.
528    ///
529    /// The downside to this is that building the record requires creating a
530    /// copy of the values, which can become slowly more expensive if the
531    /// dictionary grows.
532    ///
533    /// Additionally, if record batches from multiple different dictionary
534    /// builders for the same column are fed into a single ipc writer, beware
535    /// that entire dictionaries are likely to be re-sent frequently even when
536    /// the majority of the values are not used by the current record batch.
537    pub fn finish_preserve_values(&mut self) -> DictionaryArray<K> {
538        let values = self.values_builder.finish_cloned();
539        let keys = self.keys_builder.finish();
540
541        let data_type = DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(T::DATA_TYPE));
542
543        let builder = keys
544            .into_data()
545            .into_builder()
546            .data_type(data_type)
547            .child_data(vec![values.into_data()]);
548
549        // SAFETY: builder is constructed from valid key/value arrays produced by the builder
550        DictionaryArray::from(unsafe { builder.build_unchecked() })
551    }
552
553    /// Returns the current null buffer as a slice
554    pub fn validity_slice(&self) -> Option<&[u8]> {
555        self.keys_builder.validity_slice()
556    }
557}
558
559impl<K: ArrowDictionaryKeyType, T: ByteArrayType, V: AsRef<T::Native>> Extend<Option<V>>
560    for GenericByteDictionaryBuilder<K, T>
561{
562    #[inline]
563    fn extend<I: IntoIterator<Item = Option<V>>>(&mut self, iter: I) {
564        for v in iter {
565            self.append_option(v)
566        }
567    }
568}
569
570#[inline]
571fn resolve_key<K, T>(
572    row_idx: usize,
573    offsets: &[T::Offset],
574    raw_data: &[u8],
575    array: &GenericByteArray<T>,
576    dedup: &mut HashTable<usize>,
577    state: &ahash::RandomState,
578    storage: &mut GenericByteBuilder<T>,
579) -> Result<K::Native, ArrowError>
580where
581    K: ArrowDictionaryKeyType,
582    T: ByteArrayType,
583{
584    let start = offsets[row_idx].as_usize();
585    let end = offsets[row_idx + 1].as_usize();
586    // SAFETY: offsets are valid by GenericByteArray invariants
587    let bytes = unsafe { raw_data.get_unchecked(start..end) };
588    let hash = state.hash_one(bytes);
589    let dict_idx = *dedup
590        .entry(
591            hash,
592            |idx| bytes == get_bytes(storage, *idx),
593            |idx| state.hash_one(get_bytes(storage, *idx)),
594        )
595        .or_insert_with(|| {
596            let idx = storage.len();
597            // SAFETY: row_idx < row_count = array.len(), slot is non-null
598            storage.append_value(unsafe { array.value_unchecked(row_idx) });
599            idx
600        })
601        .get();
602    K::Native::from_usize(dict_idx).ok_or(ArrowError::DictionaryKeyOverflowError)
603}
604
605fn get_bytes<T: ByteArrayType>(values: &GenericByteBuilder<T>, idx: usize) -> &[u8] {
606    let offsets = values.offsets_slice();
607    let values = values.values_slice();
608
609    let end_offset = offsets[idx + 1].as_usize();
610    let start_offset = offsets[idx].as_usize();
611
612    &values[start_offset..end_offset]
613}
614
615/// Builder for [`DictionaryArray`] of [`StringArray`](crate::array::StringArray)
616///
617/// ```
618/// // Create a dictionary array indexed by bytes whose values are Strings.
619/// // It can thus hold up to 256 distinct string values.
620///
621/// # use arrow_array::builder::StringDictionaryBuilder;
622/// # use arrow_array::{Int8Array, StringArray};
623/// # use arrow_array::types::Int8Type;
624///
625/// let mut builder = StringDictionaryBuilder::<Int8Type>::new();
626///
627/// // The builder builds the dictionary value by value
628/// builder.append("abc").unwrap();
629/// builder.append_null();
630/// builder.append_n("def", 2).unwrap();  // appends "def" twice with a single lookup
631/// builder.append("abc").unwrap();
632/// let array = builder.finish();
633///
634/// assert_eq!(
635///   array.keys(),
636///   &Int8Array::from(vec![Some(0), None, Some(1), Some(1), Some(0)])
637/// );
638///
639/// // Values are polymorphic and so require a downcast.
640/// let av = array.values();
641/// let ava: &StringArray = av.as_any().downcast_ref::<StringArray>().unwrap();
642///
643/// assert_eq!(ava.value(0), "abc");
644/// assert_eq!(ava.value(1), "def");
645///
646/// ```
647pub type StringDictionaryBuilder<K> = GenericByteDictionaryBuilder<K, GenericStringType<i32>>;
648
649/// Builder for [`DictionaryArray`] of [`LargeStringArray`](crate::array::LargeStringArray)
650pub type LargeStringDictionaryBuilder<K> = GenericByteDictionaryBuilder<K, GenericStringType<i64>>;
651
652/// Builder for [`DictionaryArray`] of [`BinaryArray`](crate::array::BinaryArray)
653///
654/// ```
655/// // Create a dictionary array indexed by bytes whose values are binary.
656/// // It can thus hold up to 256 distinct binary values.
657///
658/// # use arrow_array::builder::BinaryDictionaryBuilder;
659/// # use arrow_array::{BinaryArray, Int8Array};
660/// # use arrow_array::types::Int8Type;
661///
662/// let mut builder = BinaryDictionaryBuilder::<Int8Type>::new();
663///
664/// // The builder builds the dictionary value by value
665/// builder.append(b"abc").unwrap();
666/// builder.append_null();
667/// builder.append(b"def").unwrap();
668/// builder.append(b"def").unwrap();
669/// builder.append(b"abc").unwrap();
670/// let array = builder.finish();
671///
672/// assert_eq!(
673///   array.keys(),
674///   &Int8Array::from(vec![Some(0), None, Some(1), Some(1), Some(0)])
675/// );
676///
677/// // Values are polymorphic and so require a downcast.
678/// let av = array.values();
679/// let ava: &BinaryArray = av.as_any().downcast_ref::<BinaryArray>().unwrap();
680///
681/// assert_eq!(ava.value(0), b"abc");
682/// assert_eq!(ava.value(1), b"def");
683///
684/// ```
685pub type BinaryDictionaryBuilder<K> = GenericByteDictionaryBuilder<K, GenericBinaryType<i32>>;
686
687/// Builder for [`DictionaryArray`] of [`LargeBinaryArray`](crate::array::LargeBinaryArray)
688pub type LargeBinaryDictionaryBuilder<K> = GenericByteDictionaryBuilder<K, GenericBinaryType<i64>>;
689
690#[cfg(test)]
691mod tests {
692    use super::*;
693
694    use crate::array::Int8Array;
695    use crate::cast::AsArray;
696    use crate::types::{
697        BinaryType, Int8Type, Int16Type, Int32Type, UInt8Type, UInt16Type, Utf8Type,
698    };
699    use crate::{ArrowPrimitiveType, BinaryArray, StringArray};
700
701    fn test_bytes_dictionary_builder<T>(values: Vec<&T::Native>)
702    where
703        T: ByteArrayType,
704        <T as ByteArrayType>::Native: PartialEq,
705        <T as ByteArrayType>::Native: AsRef<<T as ByteArrayType>::Native>,
706    {
707        let mut builder = GenericByteDictionaryBuilder::<Int8Type, T>::new();
708        builder.append(values[0]).unwrap();
709        builder.append_null();
710        builder.append(values[1]).unwrap();
711        builder.append(values[1]).unwrap();
712        builder.append(values[0]).unwrap();
713        let array = builder.finish();
714
715        assert_eq!(
716            array.keys(),
717            &Int8Array::from(vec![Some(0), None, Some(1), Some(1), Some(0)])
718        );
719
720        // Values are polymorphic and so require a downcast.
721        let av = array.values();
722        let ava: &GenericByteArray<T> = av.as_any().downcast_ref::<GenericByteArray<T>>().unwrap();
723
724        assert_eq!(*ava.value(0), *values[0]);
725        assert_eq!(*ava.value(1), *values[1]);
726    }
727
728    #[test]
729    fn test_string_dictionary_builder() {
730        test_bytes_dictionary_builder::<GenericStringType<i32>>(vec!["abc", "def"]);
731    }
732
733    #[test]
734    fn test_binary_dictionary_builder() {
735        test_bytes_dictionary_builder::<GenericBinaryType<i32>>(vec![b"abc", b"def"]);
736    }
737
738    #[test]
739    fn test_with_capacity_presizes_dedup() {
740        // `with_capacity` must size the dedup `HashTable` from `value_capacity`,
741        // otherwise the first inserts force a chain of resize+rehash cycles.
742        let value_capacity = 128;
743        let builder =
744            GenericByteDictionaryBuilder::<Int32Type, GenericStringType<i32>>::with_capacity(
745                256,
746                value_capacity,
747                value_capacity * 32,
748            );
749        assert!(
750            builder.dedup.capacity() >= value_capacity,
751            "dedup HashTable not pre-sized: got capacity {}, expected >= {}",
752            builder.dedup.capacity(),
753            value_capacity,
754        );
755    }
756
757    fn test_bytes_dictionary_builder_finish_cloned<T>(values: Vec<&T::Native>)
758    where
759        T: ByteArrayType,
760        <T as ByteArrayType>::Native: PartialEq,
761        <T as ByteArrayType>::Native: AsRef<<T as ByteArrayType>::Native>,
762    {
763        let mut builder = GenericByteDictionaryBuilder::<Int8Type, T>::new();
764
765        builder.append(values[0]).unwrap();
766        builder.append_null();
767        builder.append(values[1]).unwrap();
768        builder.append(values[1]).unwrap();
769        builder.append(values[0]).unwrap();
770        let mut array = builder.finish_cloned();
771
772        assert_eq!(
773            array.keys(),
774            &Int8Array::from(vec![Some(0), None, Some(1), Some(1), Some(0)])
775        );
776
777        // Values are polymorphic and so require a downcast.
778        let av = array.values();
779        let ava: &GenericByteArray<T> = av.as_any().downcast_ref::<GenericByteArray<T>>().unwrap();
780
781        assert_eq!(ava.value(0), values[0]);
782        assert_eq!(ava.value(1), values[1]);
783
784        builder.append(values[0]).unwrap();
785        builder.append(values[2]).unwrap();
786        builder.append(values[1]).unwrap();
787
788        array = builder.finish();
789
790        assert_eq!(
791            array.keys(),
792            &Int8Array::from(vec![
793                Some(0),
794                None,
795                Some(1),
796                Some(1),
797                Some(0),
798                Some(0),
799                Some(2),
800                Some(1)
801            ])
802        );
803
804        // Values are polymorphic and so require a downcast.
805        let av2 = array.values();
806        let ava2: &GenericByteArray<T> =
807            av2.as_any().downcast_ref::<GenericByteArray<T>>().unwrap();
808
809        assert_eq!(ava2.value(0), values[0]);
810        assert_eq!(ava2.value(1), values[1]);
811        assert_eq!(ava2.value(2), values[2]);
812    }
813
814    #[test]
815    fn test_string_dictionary_builder_finish_cloned() {
816        test_bytes_dictionary_builder_finish_cloned::<GenericStringType<i32>>(vec![
817            "abc", "def", "ghi",
818        ]);
819    }
820
821    #[test]
822    fn test_binary_dictionary_builder_finish_cloned() {
823        test_bytes_dictionary_builder_finish_cloned::<GenericBinaryType<i32>>(vec![
824            b"abc", b"def", b"ghi",
825        ]);
826    }
827
828    fn _test_try_new_from_builder_generic_for_key_types<K1, K2, T>(values: Vec<&T::Native>)
829    where
830        K1: ArrowDictionaryKeyType,
831        K1::Native: NumCast,
832        K2: ArrowDictionaryKeyType,
833        K2::Native: NumCast + From<u8>,
834        T: ByteArrayType,
835        <T as ByteArrayType>::Native: PartialEq + AsRef<<T as ByteArrayType>::Native>,
836    {
837        let mut source = GenericByteDictionaryBuilder::<K1, T>::new();
838        source.append(values[0]).unwrap();
839        source.append(values[1]).unwrap();
840        source.append_null();
841        source.append(values[2]).unwrap();
842
843        let mut result =
844            GenericByteDictionaryBuilder::<K2, T>::try_new_from_builder(source).unwrap();
845        let array = result.finish();
846
847        let mut expected_keys_builder = PrimitiveBuilder::<K2>::new();
848        expected_keys_builder
849            .append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(0u8));
850        expected_keys_builder
851            .append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(1u8));
852        expected_keys_builder.append_null();
853        expected_keys_builder
854            .append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(2u8));
855        let expected_keys = expected_keys_builder.finish();
856        assert_eq!(array.keys(), &expected_keys);
857
858        let av = array.values();
859        let ava: &GenericByteArray<T> = av.as_any().downcast_ref::<GenericByteArray<T>>().unwrap();
860        assert_eq!(ava.value(0), values[0]);
861        assert_eq!(ava.value(1), values[1]);
862        assert_eq!(ava.value(2), values[2]);
863    }
864
865    fn test_try_new_from_builder<T>(values: Vec<&T::Native>)
866    where
867        T: ByteArrayType,
868        <T as ByteArrayType>::Native: PartialEq + AsRef<<T as ByteArrayType>::Native>,
869    {
870        // test cast to bigger size unsigned
871        _test_try_new_from_builder_generic_for_key_types::<UInt8Type, UInt16Type, T>(
872            values.clone(),
873        );
874        // test cast going to smaller size unsigned
875        _test_try_new_from_builder_generic_for_key_types::<UInt16Type, UInt8Type, T>(
876            values.clone(),
877        );
878        // test cast going to bigger size signed
879        _test_try_new_from_builder_generic_for_key_types::<Int8Type, Int16Type, T>(values.clone());
880        // test cast going to smaller size signed
881        _test_try_new_from_builder_generic_for_key_types::<Int32Type, Int16Type, T>(values.clone());
882        // test going from signed to signed for different size changes
883        _test_try_new_from_builder_generic_for_key_types::<UInt8Type, Int16Type, T>(values.clone());
884        _test_try_new_from_builder_generic_for_key_types::<Int8Type, UInt8Type, T>(values.clone());
885        _test_try_new_from_builder_generic_for_key_types::<Int8Type, UInt16Type, T>(values.clone());
886        _test_try_new_from_builder_generic_for_key_types::<Int32Type, Int16Type, T>(values.clone());
887    }
888
889    #[test]
890    fn test_string_dictionary_builder_try_new_from_builder() {
891        test_try_new_from_builder::<GenericStringType<i32>>(vec!["abc", "def", "ghi"]);
892    }
893
894    #[test]
895    fn test_binary_dictionary_builder_try_new_from_builder() {
896        test_try_new_from_builder::<GenericBinaryType<i32>>(vec![b"abc", b"def", b"ghi"]);
897    }
898
899    #[test]
900    fn test_try_new_from_builder_cast_fails() {
901        let mut source_builder = StringDictionaryBuilder::<UInt16Type>::new();
902        for i in 0..257 {
903            source_builder.append_value(format!("val{i}"));
904        }
905
906        // there should be too many values that we can't downcast to the underlying type
907        // we have keys that wouldn't fit into UInt8Type
908        let result = StringDictionaryBuilder::<UInt8Type>::try_new_from_builder(source_builder);
909        assert!(result.is_err());
910        if let Err(e) = result {
911            assert!(matches!(e, ArrowError::CastError(_)));
912            assert_eq!(
913                e.to_string(),
914                "Cast error: Can't cast dictionary keys from source type UInt16 to type UInt8"
915            );
916        }
917    }
918
919    fn test_bytes_dictionary_builder_with_existing_dictionary<T>(
920        dictionary: GenericByteArray<T>,
921        values: Vec<&T::Native>,
922    ) where
923        T: ByteArrayType,
924        <T as ByteArrayType>::Native: PartialEq,
925        <T as ByteArrayType>::Native: AsRef<<T as ByteArrayType>::Native>,
926    {
927        let mut builder =
928            GenericByteDictionaryBuilder::<Int8Type, T>::new_with_dictionary(6, &dictionary)
929                .unwrap();
930        builder.append(values[0]).unwrap();
931        builder.append_null();
932        builder.append(values[1]).unwrap();
933        builder.append(values[1]).unwrap();
934        builder.append(values[0]).unwrap();
935        builder.append(values[2]).unwrap();
936        let array = builder.finish();
937
938        assert_eq!(
939            array.keys(),
940            &Int8Array::from(vec![Some(2), None, Some(1), Some(1), Some(2), Some(3)])
941        );
942
943        // Values are polymorphic and so require a downcast.
944        let av = array.values();
945        let ava: &GenericByteArray<T> = av.as_any().downcast_ref::<GenericByteArray<T>>().unwrap();
946
947        assert!(!ava.is_valid(0));
948        assert_eq!(ava.value(1), values[1]);
949        assert_eq!(ava.value(2), values[0]);
950        assert_eq!(ava.value(3), values[2]);
951    }
952
953    #[test]
954    fn test_string_dictionary_builder_with_existing_dictionary() {
955        test_bytes_dictionary_builder_with_existing_dictionary::<GenericStringType<i32>>(
956            StringArray::from(vec![None, Some("def"), Some("abc")]),
957            vec!["abc", "def", "ghi"],
958        );
959    }
960
961    #[test]
962    fn test_binary_dictionary_builder_with_existing_dictionary() {
963        let values: Vec<Option<&[u8]>> = vec![None, Some(b"def"), Some(b"abc")];
964        test_bytes_dictionary_builder_with_existing_dictionary::<GenericBinaryType<i32>>(
965            BinaryArray::from(values),
966            vec![b"abc", b"def", b"ghi"],
967        );
968    }
969
970    fn test_bytes_dictionary_builder_with_reserved_null_value<T>(
971        dictionary: GenericByteArray<T>,
972        values: Vec<&T::Native>,
973    ) where
974        T: ByteArrayType,
975        <T as ByteArrayType>::Native: PartialEq,
976        <T as ByteArrayType>::Native: AsRef<<T as ByteArrayType>::Native>,
977    {
978        let mut builder =
979            GenericByteDictionaryBuilder::<Int16Type, T>::new_with_dictionary(4, &dictionary)
980                .unwrap();
981        builder.append(values[0]).unwrap();
982        builder.append_null();
983        builder.append(values[1]).unwrap();
984        builder.append(values[0]).unwrap();
985        let array = builder.finish();
986
987        assert!(array.is_null(1));
988        assert!(!array.is_valid(1));
989
990        let keys = array.keys();
991
992        assert_eq!(keys.value(0), 1);
993        assert!(keys.is_null(1));
994        // zero initialization is currently guaranteed by Buffer allocation and resizing
995        assert_eq!(keys.value(1), 0);
996        assert_eq!(keys.value(2), 2);
997        assert_eq!(keys.value(3), 1);
998    }
999
1000    #[test]
1001    fn test_string_dictionary_builder_with_reserved_null_value() {
1002        let v: Vec<Option<&str>> = vec![None];
1003        test_bytes_dictionary_builder_with_reserved_null_value::<GenericStringType<i32>>(
1004            StringArray::from(v),
1005            vec!["abc", "def"],
1006        );
1007    }
1008
1009    #[test]
1010    fn test_binary_dictionary_builder_with_reserved_null_value() {
1011        let values: Vec<Option<&[u8]>> = vec![None];
1012        test_bytes_dictionary_builder_with_reserved_null_value::<GenericBinaryType<i32>>(
1013            BinaryArray::from(values),
1014            vec![b"abc", b"def"],
1015        );
1016    }
1017
1018    #[test]
1019    fn test_extend() {
1020        let mut builder = GenericByteDictionaryBuilder::<Int32Type, Utf8Type>::new();
1021        builder.extend(["a", "b", "c", "a", "b", "c"].into_iter().map(Some));
1022        builder.extend(["c", "d", "a"].into_iter().map(Some));
1023        let dict = builder.finish();
1024        assert_eq!(dict.keys().values(), &[0, 1, 2, 0, 1, 2, 2, 3, 0]);
1025        assert_eq!(dict.values().len(), 4);
1026    }
1027
1028    #[test]
1029    fn test_extend_dictionary() {
1030        let some_dict = {
1031            let mut builder = GenericByteDictionaryBuilder::<Int32Type, Utf8Type>::new();
1032            builder.extend(["a", "b", "c", "a", "b", "c"].into_iter().map(Some));
1033            builder.extend([None::<&str>]);
1034            builder.extend(["c", "d", "a"].into_iter().map(Some));
1035            builder.append_null();
1036            builder.finish()
1037        };
1038
1039        let mut builder = GenericByteDictionaryBuilder::<Int32Type, Utf8Type>::new();
1040        builder.extend(["e", "e", "f", "e", "d"].into_iter().map(Some));
1041        builder
1042            .extend_dictionary(&some_dict.downcast_dict().unwrap())
1043            .unwrap();
1044        let dict = builder.finish();
1045
1046        assert_eq!(dict.values().len(), 6);
1047
1048        let values = dict
1049            .downcast_dict::<GenericByteArray<Utf8Type>>()
1050            .unwrap()
1051            .into_iter()
1052            .collect::<Vec<_>>();
1053
1054        assert_eq!(
1055            values,
1056            [
1057                Some("e"),
1058                Some("e"),
1059                Some("f"),
1060                Some("e"),
1061                Some("d"),
1062                Some("a"),
1063                Some("b"),
1064                Some("c"),
1065                Some("a"),
1066                Some("b"),
1067                Some("c"),
1068                None,
1069                Some("c"),
1070                Some("d"),
1071                Some("a"),
1072                None
1073            ]
1074        );
1075    }
1076    #[test]
1077    fn test_extend_dictionary_with_null_in_mapped_value() {
1078        let some_dict = {
1079            let mut values_builder = GenericByteBuilder::<Utf8Type>::new();
1080            let mut keys_builder = PrimitiveBuilder::<Int32Type>::new();
1081
1082            // Manually build a dictionary values that the mapped values have null
1083            values_builder.append_null();
1084            keys_builder.append_value(0);
1085            values_builder.append_value("I like worm hugs");
1086            keys_builder.append_value(1);
1087
1088            let values = values_builder.finish();
1089            let keys = keys_builder.finish();
1090
1091            let data_type = DataType::Dictionary(
1092                Box::new(Int32Type::DATA_TYPE),
1093                Box::new(Utf8Type::DATA_TYPE),
1094            );
1095
1096            let builder = keys
1097                .into_data()
1098                .into_builder()
1099                .data_type(data_type)
1100                .child_data(vec![values.into_data()]);
1101
1102            DictionaryArray::from(unsafe { builder.build_unchecked() })
1103        };
1104
1105        let some_dict_values = some_dict.values().as_string::<i32>();
1106        assert_eq!(
1107            some_dict_values.into_iter().collect::<Vec<_>>(),
1108            &[None, Some("I like worm hugs")]
1109        );
1110
1111        let mut builder = GenericByteDictionaryBuilder::<Int32Type, Utf8Type>::new();
1112        builder
1113            .extend_dictionary(&some_dict.downcast_dict().unwrap())
1114            .unwrap();
1115        let dict = builder.finish();
1116
1117        assert_eq!(dict.values().len(), 1);
1118
1119        let values = dict
1120            .downcast_dict::<GenericByteArray<Utf8Type>>()
1121            .unwrap()
1122            .into_iter()
1123            .collect::<Vec<_>>();
1124
1125        assert_eq!(values, [None, Some("I like worm hugs")]);
1126    }
1127
1128    #[test]
1129    fn test_extend_all_null_dictionary() {
1130        let some_dict = {
1131            let mut builder = GenericByteDictionaryBuilder::<Int32Type, Utf8Type>::new();
1132            builder.append_nulls(2);
1133            builder.finish()
1134        };
1135
1136        let mut builder = GenericByteDictionaryBuilder::<Int32Type, Utf8Type>::new();
1137        builder
1138            .extend_dictionary(&some_dict.downcast_dict().unwrap())
1139            .unwrap();
1140        let dict = builder.finish();
1141
1142        assert_eq!(dict.values().len(), 0);
1143
1144        let values = dict
1145            .downcast_dict::<GenericByteArray<Utf8Type>>()
1146            .unwrap()
1147            .into_iter()
1148            .collect::<Vec<_>>();
1149
1150        assert_eq!(values, [None, None]);
1151    }
1152
1153    #[test]
1154    fn test_finish_preserve_values() {
1155        // Create the first dictionary
1156        let mut builder = GenericByteDictionaryBuilder::<Int32Type, Utf8Type>::new();
1157        builder.append("a").unwrap();
1158        builder.append("b").unwrap();
1159        builder.append("c").unwrap();
1160        let dict = builder.finish_preserve_values();
1161        assert_eq!(dict.keys().values(), &[0, 1, 2]);
1162        assert_eq!(dict.values().len(), 3);
1163        let values = dict
1164            .downcast_dict::<GenericByteArray<Utf8Type>>()
1165            .unwrap()
1166            .into_iter()
1167            .collect::<Vec<_>>();
1168        assert_eq!(values, [Some("a"), Some("b"), Some("c")]);
1169
1170        // Create a new dictionary
1171        builder.append("d").unwrap();
1172        builder.append("e").unwrap();
1173        let dict2 = builder.finish_preserve_values();
1174
1175        // Make sure the keys are assigned after the old ones and we have the
1176        // right values
1177        assert_eq!(dict2.keys().values(), &[3, 4]);
1178        let values = dict2
1179            .downcast_dict::<GenericByteArray<Utf8Type>>()
1180            .unwrap()
1181            .into_iter()
1182            .collect::<Vec<_>>();
1183        assert_eq!(values, [Some("d"), Some("e")]);
1184
1185        // Check that we have all of the expected values
1186        assert_eq!(dict2.values().len(), 5);
1187        let all_values = dict2
1188            .values()
1189            .as_any()
1190            .downcast_ref::<StringArray>()
1191            .unwrap()
1192            .into_iter()
1193            .collect::<Vec<_>>();
1194        assert_eq!(
1195            all_values,
1196            [Some("a"), Some("b"), Some("c"), Some("d"), Some("e"),]
1197        );
1198    }
1199
1200    #[test]
1201    fn test_append_array_deduplicates_repeated_values() {
1202        let input = StringArray::from(vec!["a", "b", "a", "c", "b", "a"]);
1203
1204        let mut builder = GenericByteDictionaryBuilder::<Int8Type, Utf8Type>::new();
1205        builder.append_array(&input).unwrap();
1206        let result = builder.finish();
1207
1208        let mut expected_builder = GenericByteDictionaryBuilder::<Int8Type, Utf8Type>::new();
1209        for value in &input {
1210            expected_builder.append_option(value);
1211        }
1212        let expected = expected_builder.finish();
1213
1214        assert_eq!(result.keys().values(), expected.keys().values());
1215        assert_eq!(result.values().len(), 3);
1216    }
1217
1218    #[test]
1219    fn test_append_array_dedup_across_consecutive_calls() {
1220        let first = StringArray::from(vec!["a", "b"]);
1221        let second = StringArray::from(vec!["b", "c", "a"]);
1222
1223        let mut builder = GenericByteDictionaryBuilder::<Int8Type, Utf8Type>::new();
1224        builder.append_array(&first).unwrap();
1225        builder.append_array(&second).unwrap();
1226        let result = builder.finish();
1227
1228        assert_eq!(result.keys().values(), &[0, 1, 1, 2, 0]);
1229        assert_eq!(result.values().len(), 3);
1230    }
1231
1232    #[test]
1233    fn test_append_array_preserves_null_positions() {
1234        let input = StringArray::from(vec![Some("x"), None, Some("x"), None, Some("y")]);
1235
1236        let mut builder = GenericByteDictionaryBuilder::<Int8Type, Utf8Type>::new();
1237        builder.append_array(&input).unwrap();
1238        let result = builder.finish();
1239
1240        let mut expected_builder = GenericByteDictionaryBuilder::<Int8Type, Utf8Type>::new();
1241        for value in &input {
1242            expected_builder.append_option(value);
1243        }
1244        let expected = expected_builder.finish();
1245
1246        assert_eq!(result.keys(), expected.keys());
1247        assert_eq!(result.values().len(), 2);
1248    }
1249
1250    #[test]
1251    fn test_append_array_overflow_binary() {
1252        // Int8 keys hold at most 128 distinct values (indices 0..=127).
1253        // Inserting a 129th distinct entry must return DictionaryKeyOverflowError.
1254        let distinct_values: Vec<Option<Vec<u8>>> = (0u16..=128)
1255            .map(|n| Some(n.to_string().into_bytes()))
1256            .collect();
1257        let input = BinaryArray::from_opt_vec(
1258            distinct_values
1259                .iter()
1260                .map(|v| v.as_deref())
1261                .collect::<Vec<_>>(),
1262        );
1263
1264        let mut builder = GenericByteDictionaryBuilder::<Int8Type, BinaryType>::new();
1265        let result = builder.append_array(&input);
1266        assert!(
1267            matches!(result, Err(ArrowError::DictionaryKeyOverflowError)),
1268            "expected DictionaryKeyOverflowError, got {result:?}"
1269        );
1270    }
1271}
arrow_array/builder/generic_bytes_dictionary_builder.rs

arrow_array/builder/
generic_bytes_dictionary_builder.rs
