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 =
463            DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(values.data_type().clone()));
464
465        let builder = keys
466            .into_data()
467            .into_builder()
468            .data_type(data_type)
469            .child_data(vec![values.into_data()]);
470
471        // SAFETY: builder is constructed from valid key/value arrays produced by the builder
472        DictionaryArray::from(unsafe { builder.build_unchecked() })
473    }
474
475    /// Builds the `DictionaryArray` without resetting the values builder or
476    /// the internal de-duplication map.
477    ///
478    /// The advantage of doing this is that the values will represent the entire
479    /// set of what has been built so-far by this builder and ensures
480    /// consistency in the assignment of keys to values across multiple calls
481    /// to `finish_preserve_values`. This enables ipc writers to efficiently
482    /// emit delta dictionaries.
483    ///
484    /// The downside to this is that building the record requires creating a
485    /// copy of the values, which can become slowly more expensive if the
486    /// dictionary grows.
487    ///
488    /// Additionally, if record batches from multiple different dictionary
489    /// builders for the same column are fed into a single ipc writer, beware
490    /// that entire dictionaries are likely to be re-sent frequently even when
491    /// the majority of the values are not used by the current record batch.
492    pub fn finish_preserve_values(&mut self) -> DictionaryArray<K> {
493        let values = self.values_builder.finish_cloned();
494        let keys = self.keys_builder.finish();
495
496        let data_type =
497            DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(values.data_type().clone()));
498
499        let builder = keys
500            .into_data()
501            .into_builder()
502            .data_type(data_type)
503            .child_data(vec![values.into_data()]);
504
505        // SAFETY: builder is constructed from valid key/value arrays produced by the builder
506        DictionaryArray::from(unsafe { builder.build_unchecked() })
507    }
508
509    /// Returns the current dictionary values buffer as a slice
510    pub fn values_slice(&self) -> &[V::Native] {
511        self.values_builder.values_slice()
512    }
513
514    /// Returns the current dictionary values buffer as a mutable slice
515    pub fn values_slice_mut(&mut self) -> &mut [V::Native] {
516        self.values_builder.values_slice_mut()
517    }
518
519    /// Returns the current null buffer as a slice
520    pub fn validity_slice(&self) -> Option<&[u8]> {
521        self.keys_builder.validity_slice()
522    }
523}
524
525impl<K: ArrowDictionaryKeyType, P: ArrowPrimitiveType> Extend<Option<P::Native>>
526    for PrimitiveDictionaryBuilder<K, P>
527{
528    #[inline]
529    fn extend<T: IntoIterator<Item = Option<P::Native>>>(&mut self, iter: T) {
530        for v in iter {
531            self.append_option(v)
532        }
533    }
534}
535
536#[cfg(test)]
537mod tests {
538    use super::*;
539
540    use crate::array::{Int32Array, UInt8Array, UInt32Array};
541    use crate::builder::Decimal128Builder;
542    use crate::cast::AsArray;
543    use crate::types::{
544        Date32Type, Decimal128Type, DurationNanosecondType, Float32Type, Float64Type, Int8Type,
545        Int16Type, Int32Type, Int64Type, TimestampNanosecondType, UInt8Type, UInt16Type,
546        UInt32Type, UInt64Type,
547    };
548
549    #[test]
550    fn test_primitive_dictionary_builder() {
551        let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::with_capacity(3, 2);
552        builder.append(12345678).unwrap();
553        builder.append_null();
554        builder.append(22345678).unwrap();
555        let array = builder.finish();
556
557        assert_eq!(
558            array.keys(),
559            &UInt8Array::from(vec![Some(0), None, Some(1)])
560        );
561
562        // Values are polymorphic and so require a downcast.
563        let av = array.values();
564        let ava: &UInt32Array = av.as_any().downcast_ref::<UInt32Array>().unwrap();
565        let avs: &[u32] = ava.values();
566
567        assert!(!array.is_null(0));
568        assert!(array.is_null(1));
569        assert!(!array.is_null(2));
570
571        assert_eq!(avs, &[12345678, 22345678]);
572    }
573
574    #[test]
575    fn test_extend() {
576        let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
577        builder.extend([1, 2, 3, 1, 2, 3, 1, 2, 3].into_iter().map(Some));
578        builder.extend([4, 5, 1, 3, 1].into_iter().map(Some));
579        let dict = builder.finish();
580        assert_eq!(
581            dict.keys().values(),
582            &[0, 1, 2, 0, 1, 2, 0, 1, 2, 3, 4, 0, 2, 0]
583        );
584        assert_eq!(dict.values().len(), 5);
585    }
586
587    #[test]
588    #[should_panic(expected = "DictionaryKeyOverflowError")]
589    fn test_primitive_dictionary_overflow() {
590        let mut builder =
591            PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::with_capacity(257, 257);
592        // 256 unique keys.
593        for i in 0..256 {
594            builder.append(i + 1000).unwrap();
595        }
596        // Special error if the key overflows (256th entry)
597        builder.append(1257).unwrap();
598    }
599
600    #[test]
601    fn test_primitive_dictionary_with_builders() {
602        let keys_builder = PrimitiveBuilder::<Int32Type>::new();
603        let values_builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(1, 2));
604        let mut builder =
605            PrimitiveDictionaryBuilder::<Int32Type, Decimal128Type>::new_from_empty_builders(
606                keys_builder,
607                values_builder,
608            );
609        let expected = DataType::Dictionary(
610            Box::new(DataType::Int32),
611            Box::new(DataType::Decimal128(1, 2)),
612        );
613        assert_eq!(builder.finish_cloned().data_type(), &expected);
614        assert_eq!(builder.finish_preserve_values().data_type(), &expected);
615
616        let dict_array = builder.finish();
617        assert_eq!(dict_array.value_type(), DataType::Decimal128(1, 2));
618        assert_eq!(dict_array.data_type(), &expected);
619    }
620
621    #[test]
622    fn test_extend_dictionary() {
623        let some_dict = {
624            let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
625            builder.extend([1, 2, 3, 1, 2, 3, 1, 2, 3].into_iter().map(Some));
626            builder.extend([None::<i32>]);
627            builder.extend([4, 5, 1, 3, 1].into_iter().map(Some));
628            builder.append_null();
629            builder.finish()
630        };
631
632        let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
633        builder.extend([6, 6, 7, 6, 5].into_iter().map(Some));
634        builder
635            .extend_dictionary(&some_dict.downcast_dict().unwrap())
636            .unwrap();
637        let dict = builder.finish();
638
639        assert_eq!(dict.values().len(), 7);
640
641        let values = dict
642            .downcast_dict::<Int32Array>()
643            .unwrap()
644            .into_iter()
645            .collect::<Vec<_>>();
646
647        assert_eq!(
648            values,
649            [
650                Some(6),
651                Some(6),
652                Some(7),
653                Some(6),
654                Some(5),
655                Some(1),
656                Some(2),
657                Some(3),
658                Some(1),
659                Some(2),
660                Some(3),
661                Some(1),
662                Some(2),
663                Some(3),
664                None,
665                Some(4),
666                Some(5),
667                Some(1),
668                Some(3),
669                Some(1),
670                None
671            ]
672        );
673    }
674
675    #[test]
676    fn test_extend_dictionary_with_null_in_mapped_value() {
677        let some_dict = {
678            let mut values_builder = PrimitiveBuilder::<Int32Type>::new();
679            let mut keys_builder = PrimitiveBuilder::<Int32Type>::new();
680
681            // Manually build a dictionary values that the mapped values have null
682            values_builder.append_null();
683            keys_builder.append_value(0);
684            values_builder.append_value(42);
685            keys_builder.append_value(1);
686
687            let values = values_builder.finish();
688            let keys = keys_builder.finish();
689
690            let data_type = DataType::Dictionary(
691                Box::new(Int32Type::DATA_TYPE),
692                Box::new(values.data_type().clone()),
693            );
694
695            let builder = keys
696                .into_data()
697                .into_builder()
698                .data_type(data_type)
699                .child_data(vec![values.into_data()]);
700
701            DictionaryArray::from(unsafe { builder.build_unchecked() })
702        };
703
704        let some_dict_values = some_dict.values().as_primitive::<Int32Type>();
705        assert_eq!(
706            some_dict_values.into_iter().collect::<Vec<_>>(),
707            &[None, Some(42)]
708        );
709
710        let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
711        builder
712            .extend_dictionary(&some_dict.downcast_dict().unwrap())
713            .unwrap();
714        let dict = builder.finish();
715
716        assert_eq!(dict.values().len(), 1);
717
718        let values = dict
719            .downcast_dict::<Int32Array>()
720            .unwrap()
721            .into_iter()
722            .collect::<Vec<_>>();
723
724        assert_eq!(values, [None, Some(42)]);
725    }
726
727    #[test]
728    fn test_extend_all_null_dictionary() {
729        let some_dict = {
730            let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
731            builder.append_nulls(2);
732            builder.finish()
733        };
734
735        let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
736        builder
737            .extend_dictionary(&some_dict.downcast_dict().unwrap())
738            .unwrap();
739        let dict = builder.finish();
740
741        assert_eq!(dict.values().len(), 0);
742
743        let values = dict
744            .downcast_dict::<Int32Array>()
745            .unwrap()
746            .into_iter()
747            .collect::<Vec<_>>();
748
749        assert_eq!(values, [None, None]);
750    }
751
752    #[test]
753    fn creating_dictionary_from_builders_should_use_values_capacity_for_the_map() {
754        let builder = PrimitiveDictionaryBuilder::<Int32Type, crate::types::TimestampMicrosecondType>::new_from_empty_builders(
755                  PrimitiveBuilder::with_capacity(1).with_data_type(DataType::Int32),
756                  PrimitiveBuilder::with_capacity(2).with_data_type(DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, Some("+08:00".into()))),
757              );
758
759        assert!(
760            builder.map.capacity() >= builder.values_builder.capacity(),
761            "map capacity {} should be at least the values capacity {}",
762            builder.map.capacity(),
763            builder.values_builder.capacity()
764        )
765    }
766
767    fn _test_try_new_from_builder_generic_for_key_types<K1, K2, V>(values: Vec<V::Native>)
768    where
769        K1: ArrowDictionaryKeyType,
770        K1::Native: NumCast,
771        K2: ArrowDictionaryKeyType,
772        K2::Native: NumCast + From<u8>,
773        V: ArrowPrimitiveType,
774    {
775        let mut source = PrimitiveDictionaryBuilder::<K1, V>::new();
776        source.append(values[0]).unwrap();
777        source.append_null();
778        source.append(values[1]).unwrap();
779        source.append(values[2]).unwrap();
780
781        let mut result = PrimitiveDictionaryBuilder::<K2, V>::try_new_from_builder(source).unwrap();
782        let array = result.finish();
783
784        let mut expected_keys_builder = PrimitiveBuilder::<K2>::new();
785        expected_keys_builder
786            .append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(0u8));
787        expected_keys_builder.append_null();
788        expected_keys_builder
789            .append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(1u8));
790        expected_keys_builder
791            .append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(2u8));
792        let expected_keys = expected_keys_builder.finish();
793        assert_eq!(array.keys(), &expected_keys);
794
795        let av = array.values();
796        let ava = av.as_any().downcast_ref::<PrimitiveArray<V>>().unwrap();
797        assert_eq!(ava.value(0), values[0]);
798        assert_eq!(ava.value(1), values[1]);
799        assert_eq!(ava.value(2), values[2]);
800    }
801
802    fn _test_try_new_from_builder_generic_for_value<T>(values: Vec<T::Native>)
803    where
804        T: ArrowPrimitiveType,
805    {
806        // test cast to bigger size unsigned
807        _test_try_new_from_builder_generic_for_key_types::<UInt8Type, UInt16Type, T>(
808            values.clone(),
809        );
810        // test cast going to smaller size unsigned
811        _test_try_new_from_builder_generic_for_key_types::<UInt16Type, UInt8Type, T>(
812            values.clone(),
813        );
814        // test cast going to bigger size signed
815        _test_try_new_from_builder_generic_for_key_types::<Int8Type, Int16Type, T>(values.clone());
816        // test cast going to smaller size signed
817        _test_try_new_from_builder_generic_for_key_types::<Int32Type, Int16Type, T>(values.clone());
818        // test going from signed to signed for different size changes
819        _test_try_new_from_builder_generic_for_key_types::<UInt8Type, Int16Type, T>(values.clone());
820        _test_try_new_from_builder_generic_for_key_types::<Int8Type, UInt8Type, T>(values.clone());
821        _test_try_new_from_builder_generic_for_key_types::<Int8Type, UInt16Type, T>(values.clone());
822        _test_try_new_from_builder_generic_for_key_types::<Int32Type, Int16Type, T>(values.clone());
823    }
824
825    #[test]
826    fn test_try_new_from_builder() {
827        // test unsigned types
828        _test_try_new_from_builder_generic_for_value::<UInt8Type>(vec![1, 2, 3]);
829        _test_try_new_from_builder_generic_for_value::<UInt16Type>(vec![1, 2, 3]);
830        _test_try_new_from_builder_generic_for_value::<UInt32Type>(vec![1, 2, 3]);
831        _test_try_new_from_builder_generic_for_value::<UInt64Type>(vec![1, 2, 3]);
832        // test signed types
833        _test_try_new_from_builder_generic_for_value::<Int8Type>(vec![-1, 0, 1]);
834        _test_try_new_from_builder_generic_for_value::<Int16Type>(vec![-1, 0, 1]);
835        _test_try_new_from_builder_generic_for_value::<Int32Type>(vec![-1, 0, 1]);
836        _test_try_new_from_builder_generic_for_value::<Int64Type>(vec![-1, 0, 1]);
837        // test some date types
838        _test_try_new_from_builder_generic_for_value::<Date32Type>(vec![5, 6, 7]);
839        _test_try_new_from_builder_generic_for_value::<DurationNanosecondType>(vec![1, 2, 3]);
840        _test_try_new_from_builder_generic_for_value::<TimestampNanosecondType>(vec![1, 2, 3]);
841        // test some floating point types
842        _test_try_new_from_builder_generic_for_value::<Float32Type>(vec![0.1, 0.2, 0.3]);
843        _test_try_new_from_builder_generic_for_value::<Float64Type>(vec![-0.1, 0.2, 0.3]);
844    }
845
846    #[test]
847    fn test_try_new_from_builder_cast_fails() {
848        let mut source_builder = PrimitiveDictionaryBuilder::<UInt16Type, UInt64Type>::new();
849        for i in 0..257 {
850            source_builder.append_value(i);
851        }
852
853        // there should be too many values that we can't downcast to the underlying type
854        // we have keys that wouldn't fit into UInt8Type
855        let result = PrimitiveDictionaryBuilder::<UInt8Type, UInt64Type>::try_new_from_builder(
856            source_builder,
857        );
858        assert!(result.is_err());
859        if let Err(e) = result {
860            assert!(matches!(e, ArrowError::CastError(_)));
861            assert_eq!(
862                e.to_string(),
863                "Cast error: Can't cast dictionary keys from source type UInt16 to type UInt8"
864            );
865        }
866    }
867
868    #[test]
869    fn test_finish_preserve_values() {
870        // Create the first dictionary
871        let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::new();
872        builder.append(10).unwrap();
873        builder.append(20).unwrap();
874        let array = builder.finish_preserve_values();
875        assert_eq!(array.keys(), &UInt8Array::from(vec![Some(0), Some(1)]));
876        let values: &[u32] = array
877            .values()
878            .as_any()
879            .downcast_ref::<UInt32Array>()
880            .unwrap()
881            .values();
882        assert_eq!(values, &[10, 20]);
883
884        // Create a new dictionary
885        builder.append(30).unwrap();
886        builder.append(40).unwrap();
887        let array2 = builder.finish_preserve_values();
888
889        // Make sure the keys are assigned after the old ones
890        // and that we have the right values
891        assert_eq!(array2.keys(), &UInt8Array::from(vec![Some(2), Some(3)]));
892        let values = array2
893            .downcast_dict::<UInt32Array>()
894            .unwrap()
895            .into_iter()
896            .collect::<Vec<_>>();
897        assert_eq!(values, vec![Some(30), Some(40)]);
898
899        // Check that we have all of the expected values
900        let all_values: &[u32] = array2
901            .values()
902            .as_any()
903            .downcast_ref::<UInt32Array>()
904            .unwrap()
905            .values();
906        assert_eq!(all_values, &[10, 20, 30, 40]);
907    }
908}
arrow_array/builder/primitive_dictionary_builder.rs

arrow_array/builder/
primitive_dictionary_builder.rs
