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

arrow_array/builder/
primitive_dictionary_builder.rs
