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

arrow_array/builder/
generic_bytes_dictionary_builder.rs
