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

arrow_array/builder/
primitive_dictionary_builder.rs
