arrow_array/builder/
primitive_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, PrimitiveBuilder};
6use crate::types::ArrowDictionaryKeyType;
7use crate::{
8    Array, ArrayRef, ArrowPrimitiveType, DictionaryArray, PrimitiveArray, TypedDictionaryArray,
9};
10use arrow_buffer::{ArrowNativeType, ToByteSlice};
11use arrow_schema::{ArrowError, DataType};
12use num_traits::NumCast;
13use std::any::Any;
14use std::collections::HashMap;
15use std::collections::hash_map::Entry;
16use std::sync::Arc;
17
18/// Wraps a type implementing `ToByteSlice` implementing `Hash` and `Eq` for it
19///
20/// This is necessary to handle types such as f32, which don't natively implement these
21#[derive(Debug)]
22struct Value<T>(T);
23
24impl<T: ToByteSlice> std::hash::Hash for Value<T> {
25    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
26        state.write(self.0.to_byte_slice())
27    }
28}
29
30impl<T: ToByteSlice> PartialEq for Value<T> {
31    fn eq(&self, other: &Self) -> bool {
32        self.0.to_byte_slice().eq(other.0.to_byte_slice())
33    }
34}
35
36impl<T: ToByteSlice> Eq for Value<T> {}
37
38/// Builder for [`DictionaryArray`] of [`PrimitiveArray`]
39///
40/// # Example:
41///
42/// ```
43///
44/// # use arrow_array::builder::PrimitiveDictionaryBuilder;
45/// # use arrow_array::types::{UInt32Type, UInt8Type};
46/// # use arrow_array::{Array, UInt32Array, UInt8Array};
47///
48/// let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::new();
49///  builder.append(12345678).unwrap();
50///  builder.append_null();
51///  builder.append(22345678).unwrap();
52///  let array = builder.finish();
53///
54///  assert_eq!(
55///      array.keys(),
56///      &UInt8Array::from(vec![Some(0), None, Some(1)])
57///  );
58///
59///  // Values are polymorphic and so require a downcast.
60///  let av = array.values();
61///  let ava: &UInt32Array = av.as_any().downcast_ref::<UInt32Array>().unwrap();
62///  let avs: &[u32] = ava.values();
63///
64///  assert!(!array.is_null(0));
65///  assert!(array.is_null(1));
66///  assert!(!array.is_null(2));
67///
68///  assert_eq!(avs, &[12345678, 22345678]);
69/// ```
70#[derive(Debug)]
71pub struct PrimitiveDictionaryBuilder<K, V>
72where
73    K: ArrowPrimitiveType,
74    V: ArrowPrimitiveType,
75{
76    keys_builder: PrimitiveBuilder<K>,
77    values_builder: PrimitiveBuilder<V>,
78    map: HashMap<Value<V::Native>, usize, ahash::RandomState>,
79}
80
81impl<K, V> Default for PrimitiveDictionaryBuilder<K, V>
82where
83    K: ArrowPrimitiveType,
84    V: ArrowPrimitiveType,
85{
86    fn default() -> Self {
87        Self::new()
88    }
89}
90
91impl<K, V> PrimitiveDictionaryBuilder<K, V>
92where
93    K: ArrowPrimitiveType,
94    V: ArrowPrimitiveType,
95{
96    /// Creates a new `PrimitiveDictionaryBuilder`.
97    pub fn new() -> Self {
98        Self {
99            keys_builder: PrimitiveBuilder::new(),
100            values_builder: PrimitiveBuilder::new(),
101            map: HashMap::default(),
102        }
103    }
104
105    /// Creates a new `PrimitiveDictionaryBuilder` from the provided keys and values builders.
106    ///
107    /// # Panics
108    ///
109    /// This method panics if `keys_builder` or `values_builder` is not empty.
110    pub fn new_from_empty_builders(
111        keys_builder: PrimitiveBuilder<K>,
112        values_builder: PrimitiveBuilder<V>,
113    ) -> Self {
114        assert!(
115            keys_builder.is_empty() && values_builder.is_empty(),
116            "keys and values builders must be empty"
117        );
118        let values_capacity = values_builder.capacity();
119        Self {
120            keys_builder,
121            values_builder,
122            map: HashMap::with_capacity_and_hasher(values_capacity, Default::default()),
123        }
124    }
125
126    /// Creates a new `PrimitiveDictionaryBuilder` from existing `PrimitiveBuilder`s of keys and values.
127    ///
128    /// # Safety
129    ///
130    /// caller must ensure that the passed in builders are valid for DictionaryArray.
131    pub unsafe fn new_from_builders(
132        keys_builder: PrimitiveBuilder<K>,
133        values_builder: PrimitiveBuilder<V>,
134    ) -> Self {
135        let keys = keys_builder.values_slice();
136        let values = values_builder.values_slice();
137        let mut map = HashMap::with_capacity_and_hasher(values.len(), Default::default());
138
139        keys.iter().zip(values.iter()).for_each(|(key, value)| {
140            map.insert(Value(*value), K::Native::to_usize(*key).unwrap());
141        });
142
143        Self {
144            keys_builder,
145            values_builder,
146            map,
147        }
148    }
149
150    /// Creates a new `PrimitiveDictionaryBuilder` with the provided capacities
151    ///
152    /// `keys_capacity`: the number of keys, i.e. length of array to build
153    /// `values_capacity`: the number of distinct dictionary values, i.e. size of dictionary
154    pub fn with_capacity(keys_capacity: usize, values_capacity: usize) -> Self {
155        Self {
156            keys_builder: PrimitiveBuilder::with_capacity(keys_capacity),
157            values_builder: PrimitiveBuilder::with_capacity(values_capacity),
158            map: HashMap::with_capacity_and_hasher(values_capacity, Default::default()),
159        }
160    }
161
162    /// Creates a new `PrimitiveDictionaryBuilder` from the existing builder with the same
163    /// keys and values, but with a new data type for the keys.
164    ///
165    /// # Example
166    /// ```
167    /// #
168    /// # use arrow_array::builder::PrimitiveDictionaryBuilder;
169    /// # use arrow_array::types::{UInt8Type, UInt16Type, UInt64Type};
170    /// # use arrow_array::UInt16Array;
171    /// # use arrow_schema::ArrowError;
172    ///
173    /// let mut u8_keyed_builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt64Type>::new();
174    ///
175    /// // appending too many values causes the dictionary to overflow
176    /// for i in 0..256 {
177    ///     u8_keyed_builder.append_value(i);
178    /// }
179    /// let result = u8_keyed_builder.append(256);
180    /// assert!(matches!(result, Err(ArrowError::DictionaryKeyOverflowError{})));
181    ///
182    /// // we need to upgrade to a larger key type
183    /// let mut u16_keyed_builder = PrimitiveDictionaryBuilder::<UInt16Type, UInt64Type>::try_new_from_builder(u8_keyed_builder).unwrap();
184    /// let dictionary_array = u16_keyed_builder.finish();
185    /// let keys = dictionary_array.keys();
186    ///
187    /// assert_eq!(keys, &UInt16Array::from_iter(0..256));
188    /// ```
189    pub fn try_new_from_builder<K2>(
190        mut source: PrimitiveDictionaryBuilder<K2, V>,
191    ) -> Result<Self, ArrowError>
192    where
193        K::Native: NumCast,
194        K2: ArrowDictionaryKeyType,
195        K2::Native: NumCast,
196    {
197        let map = source.map;
198        let values_builder = source.values_builder;
199
200        let source_keys = source.keys_builder.finish();
201        let new_keys: PrimitiveArray<K> = source_keys.try_unary(|value| {
202            num_traits::cast::cast::<K2::Native, K::Native>(value).ok_or_else(|| {
203                ArrowError::CastError(format!(
204                    "Can't cast dictionary keys from source type {:?} to type {:?}",
205                    K2::DATA_TYPE,
206                    K::DATA_TYPE
207                ))
208            })
209        })?;
210
211        // drop source key here because currently source_keys and new_keys are holding reference to
212        // the same underlying null_buffer. Below we want to call new_keys.into_builder() it must
213        // be the only reference holder.
214        drop(source_keys);
215
216        Ok(Self {
217            map,
218            keys_builder: new_keys.into_builder().map_err(|_| {
219                ArrowError::ComputeError(
220                    "Internal Error: the keys just derived from the source builder are \
221                     unexpectedly shared, so they cannot be reused as a builder"
222                        .to_string(),
223                )
224            })?,
225            values_builder,
226        })
227    }
228}
229
230impl<K, V> ArrayBuilder for PrimitiveDictionaryBuilder<K, V>
231where
232    K: ArrowDictionaryKeyType,
233    V: ArrowPrimitiveType,
234{
235    /// Returns the builder as an non-mutable `Any` reference.
236    fn as_any(&self) -> &dyn Any {
237        self
238    }
239
240    /// Returns the builder as an mutable `Any` reference.
241    fn as_any_mut(&mut self) -> &mut dyn Any {
242        self
243    }
244
245    /// Returns the boxed builder as a box of `Any`.
246    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
247        self
248    }
249
250    /// Returns the number of array slots in the builder
251    fn len(&self) -> usize {
252        self.keys_builder.len()
253    }
254
255    /// Builds the array and reset this builder.
256    fn finish(&mut self) -> ArrayRef {
257        Arc::new(self.finish())
258    }
259
260    /// Builds the array without resetting the builder.
261    fn finish_cloned(&self) -> ArrayRef {
262        Arc::new(self.finish_cloned())
263    }
264
265    fn finish_preserve_values(&mut self) -> ArrayRef {
266        Arc::new(self.finish_preserve_values())
267    }
268}
269
270impl<K, V> PrimitiveDictionaryBuilder<K, V>
271where
272    K: ArrowDictionaryKeyType,
273    V: ArrowPrimitiveType,
274{
275    #[inline]
276    fn get_or_insert_key(&mut self, value: V::Native) -> Result<K::Native, ArrowError> {
277        let key = match self.map.entry(Value(value)) {
278            Entry::Occupied(entry) => *entry.get(),
279            Entry::Vacant(entry) => {
280                let key = self.values_builder.len();
281                self.values_builder.append_value(value);
282                *entry.insert(key)
283            }
284        };
285        K::Native::from_usize(key).ok_or(ArrowError::DictionaryKeyOverflowError)
286    }
287
288    /// Append a primitive 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    #[inline]
292    pub fn append(&mut self, value: V::Native) -> Result<K::Native, ArrowError> {
293        let key = self.get_or_insert_key(value)?;
294        self.keys_builder.append_value(key);
295        Ok(key)
296    }
297
298    /// Append a value multiple times to the array.
299    /// This is the same as `append` but allows to append the same value multiple times without doing multiple lookups.
300    ///
301    /// Returns an error if the new index would overflow the key type.
302    pub fn append_n(&mut self, value: V::Native, count: usize) -> Result<K::Native, ArrowError> {
303        let key = self.get_or_insert_key(value)?;
304        self.keys_builder.append_value_n(key, count);
305        Ok(key)
306    }
307
308    /// Infallibly append a value to this builder
309    ///
310    /// # Panics
311    ///
312    /// Panics if the resulting length of the dictionary values array would exceed `T::Native::MAX`
313    #[inline]
314    pub fn append_value(&mut self, value: V::Native) {
315        self.append(value).expect("dictionary key overflow");
316    }
317
318    /// Infallibly append a value to this builder repeatedly `count` times.
319    /// This is the same as `append_value` but allows to append the same value multiple times without doing multiple lookups.
320    ///
321    /// # Panics
322    ///
323    /// Panics if the resulting length of the dictionary values array would exceed `T::Native::MAX`
324    pub fn append_values(&mut self, value: V::Native, count: usize) {
325        self.append_n(value, count)
326            .expect("dictionary key overflow");
327    }
328
329    /// Appends a null slot into the builder
330    #[inline]
331    pub fn append_null(&mut self) {
332        self.keys_builder.append_null()
333    }
334
335    /// Append `n` null slots into the builder
336    #[inline]
337    pub fn append_nulls(&mut self, n: usize) {
338        self.keys_builder.append_nulls(n)
339    }
340
341    /// Append an `Option` value into the builder
342    ///
343    /// # Panics
344    ///
345    /// Panics if the resulting length of the dictionary values array would exceed `T::Native::MAX`
346    #[inline]
347    pub fn append_option(&mut self, value: Option<V::Native>) {
348        match value {
349            None => self.append_null(),
350            Some(v) => self.append_value(v),
351        }
352    }
353
354    /// Append an `Option` value into the builder repeatedly `count` times.
355    /// This is the same as `append_option` but allows to append the same value multiple times without doing multiple lookups.
356    ///
357    /// # Panics
358    ///
359    /// Panics if the resulting length of the dictionary values array would exceed `T::Native::MAX`
360    pub fn append_options(&mut self, value: Option<V::Native>, count: usize) {
361        match value {
362            None => self.keys_builder.append_nulls(count),
363            Some(v) => self.append_values(v, count),
364        }
365    }
366
367    /// Extends builder with dictionary
368    ///
369    /// This is the same as [`Self::extend`] but is faster as it translates
370    /// the dictionary values once rather than doing a lookup for each item in the iterator
371    ///
372    /// when dictionary values are null (the actual mapped values) the keys are null
373    ///
374    pub fn extend_dictionary(
375        &mut self,
376        dictionary: &TypedDictionaryArray<K, PrimitiveArray<V>>,
377    ) -> Result<(), ArrowError> {
378        let values = dictionary.values();
379
380        let v_len = values.len();
381        let k_len = dictionary.keys().len();
382        if v_len == 0 && k_len == 0 {
383            return Ok(());
384        }
385
386        // All nulls
387        if v_len == 0 {
388            self.append_nulls(k_len);
389            return Ok(());
390        }
391
392        if k_len == 0 {
393            return Err(ArrowError::InvalidArgumentError(
394                "Dictionary keys should not be empty when values are not empty".to_string(),
395            ));
396        }
397
398        // Orphan values will be carried over to the new dictionary
399        let mapped_values = values
400            .iter()
401            // Dictionary values can technically be null, so we need to handle that
402            .map(|dict_value| {
403                dict_value
404                    .map(|dict_value| self.get_or_insert_key(dict_value))
405                    .transpose()
406            })
407            .collect::<Result<Vec<_>, _>>()?;
408
409        // Just insert the keys without additional lookups
410        dictionary.keys().iter().for_each(|key| match key {
411            None => self.append_null(),
412            Some(original_dict_index) => {
413                let index = original_dict_index.as_usize().min(v_len - 1);
414                match mapped_values[index] {
415                    None => self.append_null(),
416                    Some(mapped_value) => self.keys_builder.append_value(mapped_value),
417                }
418            }
419        });
420
421        Ok(())
422    }
423
424    /// Builds the `DictionaryArray` and reset this builder.
425    pub fn finish(&mut self) -> DictionaryArray<K> {
426        self.map.clear();
427        let values = self.values_builder.finish();
428        let keys = self.keys_builder.finish();
429
430        let data_type =
431            DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(values.data_type().clone()));
432
433        let builder = keys
434            .into_data()
435            .into_builder()
436            .data_type(data_type)
437            .child_data(vec![values.into_data()]);
438
439        // SAFETY: builder is constructed from valid key/value arrays produced by the builder
440        DictionaryArray::from(unsafe { builder.build_unchecked() })
441    }
442
443    /// Builds the `DictionaryArray` without resetting the builder.
444    pub fn finish_cloned(&self) -> DictionaryArray<K> {
445        let values = self.values_builder.finish_cloned();
446        let keys = self.keys_builder.finish_cloned();
447
448        let data_type =
449            DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(values.data_type().clone()));
450
451        let builder = keys
452            .into_data()
453            .into_builder()
454            .data_type(data_type)
455            .child_data(vec![values.into_data()]);
456
457        // SAFETY: builder is constructed from valid key/value arrays produced by the builder
458        DictionaryArray::from(unsafe { builder.build_unchecked() })
459    }
460
461    /// Builds the `DictionaryArray` without resetting the values builder or
462    /// the internal de-duplication map.
463    ///
464    /// The advantage of doing this is that the values will represent the entire
465    /// set of what has been built so-far by this builder and ensures
466    /// consistency in the assignment of keys to values across multiple calls
467    /// to `finish_preserve_values`. This enables ipc writers to efficiently
468    /// emit delta dictionaries.
469    ///
470    /// The downside to this is that building the record requires creating a
471    /// copy of the values, which can become slowly more expensive if the
472    /// dictionary grows.
473    ///
474    /// Additionally, if record batches from multiple different dictionary
475    /// builders for the same column are fed into a single ipc writer, beware
476    /// that entire dictionaries are likely to be re-sent frequently even when
477    /// the majority of the values are not used by the current record batch.
478    pub fn finish_preserve_values(&mut self) -> DictionaryArray<K> {
479        let values = self.values_builder.finish_cloned();
480        let keys = self.keys_builder.finish();
481
482        let data_type =
483            DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(values.data_type().clone()));
484
485        let builder = keys
486            .into_data()
487            .into_builder()
488            .data_type(data_type)
489            .child_data(vec![values.into_data()]);
490
491        // SAFETY: builder is constructed from valid key/value arrays produced by the builder
492        DictionaryArray::from(unsafe { builder.build_unchecked() })
493    }
494
495    /// Returns the current dictionary values buffer as a slice
496    pub fn values_slice(&self) -> &[V::Native] {
497        self.values_builder.values_slice()
498    }
499
500    /// Returns the current dictionary values buffer as a mutable slice
501    pub fn values_slice_mut(&mut self) -> &mut [V::Native] {
502        self.values_builder.values_slice_mut()
503    }
504
505    /// Returns the current null buffer as a slice
506    pub fn validity_slice(&self) -> Option<&[u8]> {
507        self.keys_builder.validity_slice()
508    }
509}
510
511impl<K: ArrowDictionaryKeyType, P: ArrowPrimitiveType> Extend<Option<P::Native>>
512    for PrimitiveDictionaryBuilder<K, P>
513{
514    #[inline]
515    fn extend<T: IntoIterator<Item = Option<P::Native>>>(&mut self, iter: T) {
516        for v in iter {
517            self.append_option(v)
518        }
519    }
520}
521
522#[cfg(test)]
523mod tests {
524    use super::*;
525
526    use crate::array::{Int32Array, UInt8Array, UInt32Array};
527    use crate::builder::Decimal128Builder;
528    use crate::cast::AsArray;
529    use crate::types::{
530        Date32Type, Decimal128Type, DurationNanosecondType, Float32Type, Float64Type, Int8Type,
531        Int16Type, Int32Type, Int64Type, TimestampNanosecondType, UInt8Type, UInt16Type,
532        UInt32Type, UInt64Type,
533    };
534
535    #[test]
536    fn test_primitive_dictionary_builder() {
537        let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::with_capacity(3, 2);
538        builder.append(12345678).unwrap();
539        builder.append_null();
540        builder.append(22345678).unwrap();
541        let array = builder.finish();
542
543        assert_eq!(
544            array.keys(),
545            &UInt8Array::from(vec![Some(0), None, Some(1)])
546        );
547
548        // Values are polymorphic and so require a downcast.
549        let av = array.values();
550        let ava: &UInt32Array = av.as_any().downcast_ref::<UInt32Array>().unwrap();
551        let avs: &[u32] = ava.values();
552
553        assert!(!array.is_null(0));
554        assert!(array.is_null(1));
555        assert!(!array.is_null(2));
556
557        assert_eq!(avs, &[12345678, 22345678]);
558    }
559
560    #[test]
561    fn test_extend() {
562        let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
563        builder.extend([1, 2, 3, 1, 2, 3, 1, 2, 3].into_iter().map(Some));
564        builder.extend([4, 5, 1, 3, 1].into_iter().map(Some));
565        let dict = builder.finish();
566        assert_eq!(
567            dict.keys().values(),
568            &[0, 1, 2, 0, 1, 2, 0, 1, 2, 3, 4, 0, 2, 0]
569        );
570        assert_eq!(dict.values().len(), 5);
571    }
572
573    #[test]
574    #[should_panic(expected = "DictionaryKeyOverflowError")]
575    fn test_primitive_dictionary_overflow() {
576        let mut builder =
577            PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::with_capacity(257, 257);
578        // 256 unique keys.
579        for i in 0..256 {
580            builder.append(i + 1000).unwrap();
581        }
582        // Special error if the key overflows (256th entry)
583        builder.append(1257).unwrap();
584    }
585
586    #[test]
587    fn test_primitive_dictionary_with_builders() {
588        let keys_builder = PrimitiveBuilder::<Int32Type>::new();
589        let values_builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(1, 2));
590        let mut builder =
591            PrimitiveDictionaryBuilder::<Int32Type, Decimal128Type>::new_from_empty_builders(
592                keys_builder,
593                values_builder,
594            );
595        let expected = DataType::Dictionary(
596            Box::new(DataType::Int32),
597            Box::new(DataType::Decimal128(1, 2)),
598        );
599        assert_eq!(builder.finish_cloned().data_type(), &expected);
600        assert_eq!(builder.finish_preserve_values().data_type(), &expected);
601
602        let dict_array = builder.finish();
603        assert_eq!(dict_array.value_type(), DataType::Decimal128(1, 2));
604        assert_eq!(dict_array.data_type(), &expected);
605    }
606
607    #[test]
608    fn test_extend_dictionary() {
609        let some_dict = {
610            let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
611            builder.extend([1, 2, 3, 1, 2, 3, 1, 2, 3].into_iter().map(Some));
612            builder.extend([None::<i32>]);
613            builder.extend([4, 5, 1, 3, 1].into_iter().map(Some));
614            builder.append_null();
615            builder.finish()
616        };
617
618        let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
619        builder.extend([6, 6, 7, 6, 5].into_iter().map(Some));
620        builder
621            .extend_dictionary(&some_dict.downcast_dict().unwrap())
622            .unwrap();
623        let dict = builder.finish();
624
625        assert_eq!(dict.values().len(), 7);
626
627        let values = dict
628            .downcast_dict::<Int32Array>()
629            .unwrap()
630            .into_iter()
631            .collect::<Vec<_>>();
632
633        assert_eq!(
634            values,
635            [
636                Some(6),
637                Some(6),
638                Some(7),
639                Some(6),
640                Some(5),
641                Some(1),
642                Some(2),
643                Some(3),
644                Some(1),
645                Some(2),
646                Some(3),
647                Some(1),
648                Some(2),
649                Some(3),
650                None,
651                Some(4),
652                Some(5),
653                Some(1),
654                Some(3),
655                Some(1),
656                None
657            ]
658        );
659    }
660
661    #[test]
662    fn test_extend_dictionary_with_null_in_mapped_value() {
663        let some_dict = {
664            let mut values_builder = PrimitiveBuilder::<Int32Type>::new();
665            let mut keys_builder = PrimitiveBuilder::<Int32Type>::new();
666
667            // Manually build a dictionary values that the mapped values have null
668            values_builder.append_null();
669            keys_builder.append_value(0);
670            values_builder.append_value(42);
671            keys_builder.append_value(1);
672
673            let values = values_builder.finish();
674            let keys = keys_builder.finish();
675
676            let data_type = DataType::Dictionary(
677                Box::new(Int32Type::DATA_TYPE),
678                Box::new(values.data_type().clone()),
679            );
680
681            let builder = keys
682                .into_data()
683                .into_builder()
684                .data_type(data_type)
685                .child_data(vec![values.into_data()]);
686
687            DictionaryArray::from(unsafe { builder.build_unchecked() })
688        };
689
690        let some_dict_values = some_dict.values().as_primitive::<Int32Type>();
691        assert_eq!(
692            some_dict_values.into_iter().collect::<Vec<_>>(),
693            &[None, Some(42)]
694        );
695
696        let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
697        builder
698            .extend_dictionary(&some_dict.downcast_dict().unwrap())
699            .unwrap();
700        let dict = builder.finish();
701
702        assert_eq!(dict.values().len(), 1);
703
704        let values = dict
705            .downcast_dict::<Int32Array>()
706            .unwrap()
707            .into_iter()
708            .collect::<Vec<_>>();
709
710        assert_eq!(values, [None, Some(42)]);
711    }
712
713    #[test]
714    fn test_extend_all_null_dictionary() {
715        let some_dict = {
716            let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
717            builder.append_nulls(2);
718            builder.finish()
719        };
720
721        let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
722        builder
723            .extend_dictionary(&some_dict.downcast_dict().unwrap())
724            .unwrap();
725        let dict = builder.finish();
726
727        assert_eq!(dict.values().len(), 0);
728
729        let values = dict
730            .downcast_dict::<Int32Array>()
731            .unwrap()
732            .into_iter()
733            .collect::<Vec<_>>();
734
735        assert_eq!(values, [None, None]);
736    }
737
738    #[test]
739    fn creating_dictionary_from_builders_should_use_values_capacity_for_the_map() {
740        let builder = PrimitiveDictionaryBuilder::<Int32Type, crate::types::TimestampMicrosecondType>::new_from_empty_builders(
741                  PrimitiveBuilder::with_capacity(1).with_data_type(DataType::Int32),
742                  PrimitiveBuilder::with_capacity(2).with_data_type(DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, Some("+08:00".into()))),
743              );
744
745        assert!(
746            builder.map.capacity() >= builder.values_builder.capacity(),
747            "map capacity {} should be at least the values capacity {}",
748            builder.map.capacity(),
749            builder.values_builder.capacity()
750        )
751    }
752
753    fn _test_try_new_from_builder_generic_for_key_types<K1, K2, V>(values: Vec<V::Native>)
754    where
755        K1: ArrowDictionaryKeyType,
756        K1::Native: NumCast,
757        K2: ArrowDictionaryKeyType,
758        K2::Native: NumCast + From<u8>,
759        V: ArrowPrimitiveType,
760    {
761        let mut source = PrimitiveDictionaryBuilder::<K1, V>::new();
762        source.append(values[0]).unwrap();
763        source.append_null();
764        source.append(values[1]).unwrap();
765        source.append(values[2]).unwrap();
766
767        let mut result = PrimitiveDictionaryBuilder::<K2, V>::try_new_from_builder(source).unwrap();
768        let array = result.finish();
769
770        let mut expected_keys_builder = PrimitiveBuilder::<K2>::new();
771        expected_keys_builder
772            .append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(0u8));
773        expected_keys_builder.append_null();
774        expected_keys_builder
775            .append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(1u8));
776        expected_keys_builder
777            .append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(2u8));
778        let expected_keys = expected_keys_builder.finish();
779        assert_eq!(array.keys(), &expected_keys);
780
781        let av = array.values();
782        let ava = av.as_any().downcast_ref::<PrimitiveArray<V>>().unwrap();
783        assert_eq!(ava.value(0), values[0]);
784        assert_eq!(ava.value(1), values[1]);
785        assert_eq!(ava.value(2), values[2]);
786    }
787
788    fn _test_try_new_from_builder_generic_for_value<T>(values: Vec<T::Native>)
789    where
790        T: ArrowPrimitiveType,
791    {
792        // test cast to bigger size unsigned
793        _test_try_new_from_builder_generic_for_key_types::<UInt8Type, UInt16Type, T>(
794            values.clone(),
795        );
796        // test cast going to smaller size unsigned
797        _test_try_new_from_builder_generic_for_key_types::<UInt16Type, UInt8Type, T>(
798            values.clone(),
799        );
800        // test cast going to bigger size signed
801        _test_try_new_from_builder_generic_for_key_types::<Int8Type, Int16Type, T>(values.clone());
802        // test cast going to smaller size signed
803        _test_try_new_from_builder_generic_for_key_types::<Int32Type, Int16Type, T>(values.clone());
804        // test going from signed to signed for different size changes
805        _test_try_new_from_builder_generic_for_key_types::<UInt8Type, Int16Type, T>(values.clone());
806        _test_try_new_from_builder_generic_for_key_types::<Int8Type, UInt8Type, T>(values.clone());
807        _test_try_new_from_builder_generic_for_key_types::<Int8Type, UInt16Type, T>(values.clone());
808        _test_try_new_from_builder_generic_for_key_types::<Int32Type, Int16Type, T>(values.clone());
809    }
810
811    #[test]
812    fn test_try_new_from_builder() {
813        // test unsigned types
814        _test_try_new_from_builder_generic_for_value::<UInt8Type>(vec![1, 2, 3]);
815        _test_try_new_from_builder_generic_for_value::<UInt16Type>(vec![1, 2, 3]);
816        _test_try_new_from_builder_generic_for_value::<UInt32Type>(vec![1, 2, 3]);
817        _test_try_new_from_builder_generic_for_value::<UInt64Type>(vec![1, 2, 3]);
818        // test signed types
819        _test_try_new_from_builder_generic_for_value::<Int8Type>(vec![-1, 0, 1]);
820        _test_try_new_from_builder_generic_for_value::<Int16Type>(vec![-1, 0, 1]);
821        _test_try_new_from_builder_generic_for_value::<Int32Type>(vec![-1, 0, 1]);
822        _test_try_new_from_builder_generic_for_value::<Int64Type>(vec![-1, 0, 1]);
823        // test some date types
824        _test_try_new_from_builder_generic_for_value::<Date32Type>(vec![5, 6, 7]);
825        _test_try_new_from_builder_generic_for_value::<DurationNanosecondType>(vec![1, 2, 3]);
826        _test_try_new_from_builder_generic_for_value::<TimestampNanosecondType>(vec![1, 2, 3]);
827        // test some floating point types
828        _test_try_new_from_builder_generic_for_value::<Float32Type>(vec![0.1, 0.2, 0.3]);
829        _test_try_new_from_builder_generic_for_value::<Float64Type>(vec![-0.1, 0.2, 0.3]);
830    }
831
832    #[test]
833    fn test_try_new_from_builder_cast_fails() {
834        let mut source_builder = PrimitiveDictionaryBuilder::<UInt16Type, UInt64Type>::new();
835        for i in 0..257 {
836            source_builder.append_value(i);
837        }
838
839        // there should be too many values that we can't downcast to the underlying type
840        // we have keys that wouldn't fit into UInt8Type
841        let result = PrimitiveDictionaryBuilder::<UInt8Type, UInt64Type>::try_new_from_builder(
842            source_builder,
843        );
844        assert!(result.is_err());
845        if let Err(e) = result {
846            assert!(matches!(e, ArrowError::CastError(_)));
847            assert_eq!(
848                e.to_string(),
849                "Cast error: Can't cast dictionary keys from source type UInt16 to type UInt8"
850            );
851        }
852    }
853
854    #[test]
855    fn test_finish_preserve_values() {
856        // Create the first dictionary
857        let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::new();
858        builder.append(10).unwrap();
859        builder.append(20).unwrap();
860        let array = builder.finish_preserve_values();
861        assert_eq!(array.keys(), &UInt8Array::from(vec![Some(0), Some(1)]));
862        let values: &[u32] = array
863            .values()
864            .as_any()
865            .downcast_ref::<UInt32Array>()
866            .unwrap()
867            .values();
868        assert_eq!(values, &[10, 20]);
869
870        // Create a new dictionary
871        builder.append(30).unwrap();
872        builder.append(40).unwrap();
873        let array2 = builder.finish_preserve_values();
874
875        // Make sure the keys are assigned after the old ones
876        // and that we have the right values
877        assert_eq!(array2.keys(), &UInt8Array::from(vec![Some(2), Some(3)]));
878        let values = array2
879            .downcast_dict::<UInt32Array>()
880            .unwrap()
881            .into_iter()
882            .collect::<Vec<_>>();
883        assert_eq!(values, vec![Some(30), Some(40)]);
884
885        // Check that we have all of the expected values
886        let all_values: &[u32] = array2
887            .values()
888            .as_any()
889            .downcast_ref::<UInt32Array>()
890            .unwrap()
891            .values();
892        assert_eq!(all_values, &[10, 20, 30, 40]);
893    }
894}
arrow_array/builder/primitive_dictionary_builder.rs

arrow_array/builder/
primitive_dictionary_builder.rs
