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

arrow_array/builder/
primitive_dictionary_builder.rs
