arrow_array/array/
dictionary_array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use crate::builder::{PrimitiveDictionaryBuilder, StringDictionaryBuilder};
6use crate::cast::AsArray;
7use crate::iterator::ArrayIter;
8use crate::types::*;
9use crate::{
10    Array, ArrayAccessor, ArrayRef, ArrowNativeTypeOp, PrimitiveArray, Scalar, StringArray,
11    make_array,
12};
13use arrow_buffer::bit_util::set_bit;
14use arrow_buffer::buffer::NullBuffer;
15use arrow_buffer::{ArrowNativeType, BooleanBuffer, BooleanBufferBuilder, ScalarBuffer};
16use arrow_data::ArrayData;
17use arrow_schema::{ArrowError, DataType};
18use std::any::Any;
19use std::sync::Arc;
20
21/// A [`DictionaryArray`] indexed by `i8`
22///
23/// # Example: Using `collect`
24/// ```
25/// # use arrow_array::{Array, Int8DictionaryArray, Int8Array, StringArray};
26/// # use std::sync::Arc;
27///
28/// let array: Int8DictionaryArray = vec!["a", "a", "b", "c"].into_iter().collect();
29/// let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "b", "c"]));
30/// assert_eq!(array.keys(), &Int8Array::from(vec![0, 0, 1, 2]));
31/// assert_eq!(array.values(), &values);
32/// ```
33///
34/// See [`DictionaryArray`] for more information and examples
35pub type Int8DictionaryArray = DictionaryArray<Int8Type>;
36
37/// A [`DictionaryArray`] indexed by `i16`
38///
39/// # Example: Using `collect`
40/// ```
41/// # use arrow_array::{Array, Int16DictionaryArray, Int16Array, StringArray};
42/// # use std::sync::Arc;
43///
44/// let array: Int16DictionaryArray = vec!["a", "a", "b", "c"].into_iter().collect();
45/// let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "b", "c"]));
46/// assert_eq!(array.keys(), &Int16Array::from(vec![0, 0, 1, 2]));
47/// assert_eq!(array.values(), &values);
48/// ```
49///
50/// See [`DictionaryArray`] for more information and examples
51pub type Int16DictionaryArray = DictionaryArray<Int16Type>;
52
53/// A [`DictionaryArray`] indexed by `i32`
54///
55/// # Example: Using `collect`
56/// ```
57/// # use arrow_array::{Array, Int32DictionaryArray, Int32Array, StringArray};
58/// # use std::sync::Arc;
59///
60/// let array: Int32DictionaryArray = vec!["a", "a", "b", "c"].into_iter().collect();
61/// let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "b", "c"]));
62/// assert_eq!(array.keys(), &Int32Array::from(vec![0, 0, 1, 2]));
63/// assert_eq!(array.values(), &values);
64/// ```
65///
66/// See [`DictionaryArray`] for more information and examples
67pub type Int32DictionaryArray = DictionaryArray<Int32Type>;
68
69/// A [`DictionaryArray`] indexed by `i64`
70///
71/// # Example: Using `collect`
72/// ```
73/// # use arrow_array::{Array, Int64DictionaryArray, Int64Array, StringArray};
74/// # use std::sync::Arc;
75///
76/// let array: Int64DictionaryArray = vec!["a", "a", "b", "c"].into_iter().collect();
77/// let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "b", "c"]));
78/// assert_eq!(array.keys(), &Int64Array::from(vec![0, 0, 1, 2]));
79/// assert_eq!(array.values(), &values);
80/// ```
81///
82/// See [`DictionaryArray`] for more information and examples
83pub type Int64DictionaryArray = DictionaryArray<Int64Type>;
84
85/// A [`DictionaryArray`] indexed by `u8`
86///
87/// # Example: Using `collect`
88/// ```
89/// # use arrow_array::{Array, UInt8DictionaryArray, UInt8Array, StringArray};
90/// # use std::sync::Arc;
91///
92/// let array: UInt8DictionaryArray = vec!["a", "a", "b", "c"].into_iter().collect();
93/// let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "b", "c"]));
94/// assert_eq!(array.keys(), &UInt8Array::from(vec![0, 0, 1, 2]));
95/// assert_eq!(array.values(), &values);
96/// ```
97///
98/// See [`DictionaryArray`] for more information and examples
99pub type UInt8DictionaryArray = DictionaryArray<UInt8Type>;
100
101/// A [`DictionaryArray`] indexed by `u16`
102///
103/// # Example: Using `collect`
104/// ```
105/// # use arrow_array::{Array, UInt16DictionaryArray, UInt16Array, StringArray};
106/// # use std::sync::Arc;
107///
108/// let array: UInt16DictionaryArray = vec!["a", "a", "b", "c"].into_iter().collect();
109/// let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "b", "c"]));
110/// assert_eq!(array.keys(), &UInt16Array::from(vec![0, 0, 1, 2]));
111/// assert_eq!(array.values(), &values);
112/// ```
113///
114/// See [`DictionaryArray`] for more information and examples
115pub type UInt16DictionaryArray = DictionaryArray<UInt16Type>;
116
117/// A [`DictionaryArray`] indexed by `u32`
118///
119/// # Example: Using `collect`
120/// ```
121/// # use arrow_array::{Array, UInt32DictionaryArray, UInt32Array, StringArray};
122/// # use std::sync::Arc;
123///
124/// let array: UInt32DictionaryArray = vec!["a", "a", "b", "c"].into_iter().collect();
125/// let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "b", "c"]));
126/// assert_eq!(array.keys(), &UInt32Array::from(vec![0, 0, 1, 2]));
127/// assert_eq!(array.values(), &values);
128/// ```
129///
130/// See [`DictionaryArray`] for more information and examples
131pub type UInt32DictionaryArray = DictionaryArray<UInt32Type>;
132
133/// A [`DictionaryArray`] indexed by `u64`
134///
135/// # Example: Using `collect`
136/// ```
137/// # use arrow_array::{Array, UInt64DictionaryArray, UInt64Array, StringArray};
138/// # use std::sync::Arc;
139///
140/// let array: UInt64DictionaryArray = vec!["a", "a", "b", "c"].into_iter().collect();
141/// let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "b", "c"]));
142/// assert_eq!(array.keys(), &UInt64Array::from(vec![0, 0, 1, 2]));
143/// assert_eq!(array.values(), &values);
144/// ```
145///
146/// See [`DictionaryArray`] for more information and examples
147pub type UInt64DictionaryArray = DictionaryArray<UInt64Type>;
148
149/// An array of [dictionary encoded values](https://arrow.apache.org/docs/format/Columnar.html#dictionary-encoded-layout)
150///
151/// This is mostly used to represent strings or a limited set of primitive types as integers,
152/// for example when doing NLP analysis or representing chromosomes by name.
153///
154/// [`DictionaryArray`] are represented using a `keys` array and a
155/// `values` array, which may be different lengths. The `keys` array
156/// stores indexes in the `values` array which holds
157/// the corresponding logical value, as shown here:
158///
159/// ```text
160/// ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─
161///   ┌─────────────────┐  ┌─────────┐ │     ┌─────────────────┐
162/// │ │        A        │  │    0    │       │        A        │     values[keys[0]]
163///   ├─────────────────┤  ├─────────┤ │     ├─────────────────┤
164/// │ │        D        │  │    2    │       │        B        │     values[keys[1]]
165///   ├─────────────────┤  ├─────────┤ │     ├─────────────────┤
166/// │ │        B        │  │    2    │       │        B        │     values[keys[2]]
167///   └─────────────────┘  ├─────────┤ │     ├─────────────────┤
168/// │                      │    1    │       │        D        │     values[keys[3]]
169///                        ├─────────┤ │     ├─────────────────┤
170/// │                      │    1    │       │        D        │     values[keys[4]]
171///                        ├─────────┤ │     ├─────────────────┤
172/// │                      │    0    │       │        A        │     values[keys[5]]
173///                        └─────────┘ │     └─────────────────┘
174/// │       values            keys
175///  ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘
176///                                             Logical array
177///                                                Contents
178///           DictionaryArray
179///              length = 6
180/// ```
181///
182/// # Example: From Nullable Data
183///
184/// ```
185/// # use arrow_array::{DictionaryArray, Int8Array, types::Int8Type};
186/// let test = vec!["a", "a", "b", "c"];
187/// let array : DictionaryArray<Int8Type> = test.iter().map(|&x| if x == "b" {None} else {Some(x)}).collect();
188/// assert_eq!(array.keys(), &Int8Array::from(vec![Some(0), Some(0), None, Some(1)]));
189/// ```
190///
191/// # Example: From Non-Nullable Data
192///
193/// ```
194/// # use arrow_array::{DictionaryArray, Int8Array, types::Int8Type};
195/// let test = vec!["a", "a", "b", "c"];
196/// let array : DictionaryArray<Int8Type> = test.into_iter().collect();
197/// assert_eq!(array.keys(), &Int8Array::from(vec![0, 0, 1, 2]));
198/// ```
199///
200/// # Example: From Existing Arrays
201///
202/// ```
203/// # use std::sync::Arc;
204/// # use arrow_array::{DictionaryArray, Int8Array, StringArray, types::Int8Type};
205/// // You can form your own DictionaryArray by providing the
206/// // values (dictionary) and keys (indexes into the dictionary):
207/// let values = StringArray::from_iter_values(["a", "b", "c"]);
208/// let keys = Int8Array::from_iter_values([0, 0, 1, 2]);
209/// let array = DictionaryArray::<Int8Type>::try_new(keys, Arc::new(values)).unwrap();
210/// let expected: DictionaryArray::<Int8Type> = vec!["a", "a", "b", "c"].into_iter().collect();
211/// assert_eq!(&array, &expected);
212/// ```
213///
214/// # Example: Using Builder
215///
216/// ```
217/// # use arrow_array::{Array, StringArray};
218/// # use arrow_array::builder::StringDictionaryBuilder;
219/// # use arrow_array::types::Int32Type;
220/// let mut builder = StringDictionaryBuilder::<Int32Type>::new();
221/// builder.append_value("a");
222/// builder.append_null();
223/// builder.append_value("a");
224/// builder.append_value("b");
225/// let array = builder.finish();
226///
227/// let values: Vec<_> = array.downcast_dict::<StringArray>().unwrap().into_iter().collect();
228/// assert_eq!(&values, &[Some("a"), None, Some("a"), Some("b")]);
229/// ```
230pub struct DictionaryArray<K: ArrowDictionaryKeyType> {
231    data_type: DataType,
232
233    /// The keys of this dictionary. These are constructed from the
234    /// buffer and null bitmap of `data`.  Also, note that these do
235    /// not correspond to the true values of this array. Rather, they
236    /// map to the real values.
237    keys: PrimitiveArray<K>,
238
239    /// Array of dictionary values (can be any DataType).
240    values: ArrayRef,
241
242    /// Values are ordered.
243    is_ordered: bool,
244}
245
246impl<K: ArrowDictionaryKeyType> Clone for DictionaryArray<K> {
247    fn clone(&self) -> Self {
248        Self {
249            data_type: self.data_type.clone(),
250            keys: self.keys.clone(),
251            values: self.values.clone(),
252            is_ordered: self.is_ordered,
253        }
254    }
255}
256
257impl<K: ArrowDictionaryKeyType> DictionaryArray<K> {
258    /// Attempt to create a new DictionaryArray with a specified keys
259    /// (indexes into the dictionary) and values (dictionary)
260    /// array.
261    ///
262    /// # Panics
263    ///
264    /// Panics if [`Self::try_new`] returns an error
265    pub fn new(keys: PrimitiveArray<K>, values: ArrayRef) -> Self {
266        Self::try_new(keys, values).unwrap()
267    }
268
269    /// Attempt to create a new DictionaryArray with a specified keys
270    /// (indexes into the dictionary) and values (dictionary)
271    /// array.
272    ///
273    /// # Errors
274    ///
275    /// Returns an error if any `keys[i] >= values.len() || keys[i] < 0`
276    pub fn try_new(keys: PrimitiveArray<K>, values: ArrayRef) -> Result<Self, ArrowError> {
277        let data_type = DataType::Dictionary(
278            Box::new(keys.data_type().clone()),
279            Box::new(values.data_type().clone()),
280        );
281
282        // // we can skip the validating the keys if they are all null
283        let all_null = keys.null_count() == keys.len();
284
285        if !all_null {
286            let zero = K::Native::usize_as(0);
287            let values_len = values.len();
288
289            if let Some((idx, v)) = keys.values().iter().enumerate().find(|(idx, v)| {
290                (v.is_lt(zero) || v.as_usize() >= values_len) && keys.is_valid(*idx)
291            }) {
292                return Err(ArrowError::InvalidArgumentError(format!(
293                    "Invalid dictionary key {v:?} at index {idx}, expected 0 <= key < {values_len}",
294                )));
295            }
296        }
297
298        Ok(Self {
299            data_type,
300            keys,
301            values,
302            is_ordered: false,
303        })
304    }
305
306    /// Create a new [`Scalar`] from `value`
307    pub fn new_scalar<T: Array + 'static>(value: Scalar<T>) -> Scalar<Self> {
308        Scalar::new(Self::new(
309            PrimitiveArray::new(vec![K::Native::usize_as(0)].into(), None),
310            Arc::new(value.into_inner()),
311        ))
312    }
313
314    /// Create a new [`DictionaryArray`] without performing validation
315    ///
316    /// # Safety
317    ///
318    /// Safe provided [`Self::try_new`] would not return an error
319    pub unsafe fn new_unchecked(keys: PrimitiveArray<K>, values: ArrayRef) -> Self {
320        if cfg!(feature = "force_validate") {
321            return Self::new(keys, values);
322        }
323
324        let data_type = DataType::Dictionary(
325            Box::new(keys.data_type().clone()),
326            Box::new(values.data_type().clone()),
327        );
328
329        Self {
330            data_type,
331            keys,
332            values,
333            is_ordered: false,
334        }
335    }
336
337    /// Deconstruct this array into its constituent parts
338    pub fn into_parts(self) -> (PrimitiveArray<K>, ArrayRef) {
339        (self.keys, self.values)
340    }
341
342    /// Return an array view of the keys of this dictionary as a PrimitiveArray.
343    pub fn keys(&self) -> &PrimitiveArray<K> {
344        &self.keys
345    }
346
347    /// If `value` is present in `values` (aka the dictionary),
348    /// returns the corresponding key (index into the `values`
349    /// array). Otherwise returns `None`.
350    ///
351    /// # Panics
352    /// Panics if `values` is not a [`StringArray`].
353    pub fn lookup_key(&self, value: &str) -> Option<K::Native> {
354        let rd_buf: &StringArray = self.values.as_any().downcast_ref::<StringArray>().unwrap();
355
356        (0..rd_buf.len())
357            .position(|i| rd_buf.value(i) == value)
358            .and_then(K::Native::from_usize)
359    }
360
361    /// Returns a reference to the dictionary values array
362    pub fn values(&self) -> &ArrayRef {
363        &self.values
364    }
365
366    /// Returns a clone of the value type of this list.
367    pub fn value_type(&self) -> DataType {
368        self.values.data_type().clone()
369    }
370
371    /// The length of the dictionary is the length of the keys array.
372    pub fn len(&self) -> usize {
373        self.keys.len()
374    }
375
376    /// Whether this dictionary is empty
377    pub fn is_empty(&self) -> bool {
378        self.keys.is_empty()
379    }
380
381    /// Currently exists for compatibility purposes with Arrow IPC.
382    pub fn is_ordered(&self) -> bool {
383        self.is_ordered
384    }
385
386    /// Return an iterator over the keys (indexes into the dictionary)
387    pub fn keys_iter(&self) -> impl Iterator<Item = Option<usize>> + '_ {
388        self.keys.iter().map(|key| key.map(|k| k.as_usize()))
389    }
390
391    /// Return the value of `keys` (the dictionary key) at index `i`,
392    /// cast to `usize`, `None` if the value at `i` is `NULL`.
393    ///
394    /// # Panics
395    /// Panics if `i >= self.len()`
396    pub fn key(&self, i: usize) -> Option<usize> {
397        self.keys.is_valid(i).then(|| self.keys.value(i).as_usize())
398    }
399
400    /// Returns a zero-copy slice of this array with the indicated offset and length.
401    ///
402    /// # Panics
403    /// Panics if `offset + length > self.len()`
404    pub fn slice(&self, offset: usize, length: usize) -> Self {
405        Self {
406            data_type: self.data_type.clone(),
407            keys: self.keys.slice(offset, length),
408            values: self.values.clone(),
409            is_ordered: self.is_ordered,
410        }
411    }
412
413    /// Downcast this dictionary to a [`TypedDictionaryArray`]
414    ///
415    /// ```
416    /// use arrow_array::{Array, ArrayAccessor, DictionaryArray, StringArray, types::Int32Type};
417    ///
418    /// let orig = [Some("a"), Some("b"), None];
419    /// let dictionary = DictionaryArray::<Int32Type>::from_iter(orig);
420    /// let typed = dictionary.downcast_dict::<StringArray>().unwrap();
421    /// assert_eq!(typed.value(0), "a");
422    /// assert_eq!(typed.value(1), "b");
423    /// assert!(typed.is_null(2));
424    /// ```
425    ///
426    pub fn downcast_dict<V: 'static>(&self) -> Option<TypedDictionaryArray<'_, K, V>> {
427        let values = self.values.as_any().downcast_ref()?;
428        Some(TypedDictionaryArray {
429            dictionary: self,
430            values,
431        })
432    }
433
434    /// Returns a new dictionary with the same keys as the current instance
435    /// but with a different set of dictionary values
436    ///
437    /// This can be used to perform an operation on the values of a dictionary
438    ///
439    /// # Panics
440    ///
441    /// Panics if `values` has a length less than the current values
442    ///
443    /// ```
444    /// # use std::sync::Arc;
445    /// # use arrow_array::builder::PrimitiveDictionaryBuilder;
446    /// # use arrow_array::{Int8Array, Int64Array, ArrayAccessor};
447    /// # use arrow_array::types::{Int32Type, Int8Type};
448    ///
449    /// // Construct a Dict(Int32, Int8)
450    /// let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int8Type>::with_capacity(2, 200);
451    /// for i in 0..100 {
452    ///     builder.append(i % 2).unwrap();
453    /// }
454    ///
455    /// let dictionary = builder.finish();
456    ///
457    /// // Perform a widening cast of dictionary values
458    /// let typed_dictionary = dictionary.downcast_dict::<Int8Array>().unwrap();
459    /// let values: Int64Array = typed_dictionary.values().unary(|x| x as i64);
460    ///
461    /// // Create a Dict(Int32,
462    /// let new = dictionary.with_values(Arc::new(values));
463    ///
464    /// // Verify values are as expected
465    /// let new_typed = new.downcast_dict::<Int64Array>().unwrap();
466    /// for i in 0..100 {
467    ///     assert_eq!(new_typed.value(i), (i % 2) as i64)
468    /// }
469    /// ```
470    ///
471    pub fn with_values(&self, values: ArrayRef) -> Self {
472        assert!(values.len() >= self.values.len());
473        let data_type =
474            DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(values.data_type().clone()));
475        Self {
476            data_type,
477            keys: self.keys.clone(),
478            values,
479            is_ordered: false,
480        }
481    }
482
483    /// Returns `PrimitiveDictionaryBuilder` of this dictionary array for mutating
484    /// its keys and values if the underlying data buffer is not shared by others.
485    #[expect(clippy::result_large_err)]
486    pub fn into_primitive_dict_builder<V>(self) -> Result<PrimitiveDictionaryBuilder<K, V>, Self>
487    where
488        V: ArrowPrimitiveType,
489    {
490        if !self.value_type().is_primitive() {
491            return Err(self);
492        }
493
494        let key_array = self.keys().clone();
495        let value_array = self.values().as_primitive::<V>().clone();
496
497        drop(self.keys);
498        drop(self.values);
499
500        let key_builder = key_array.into_builder();
501        let value_builder = value_array.into_builder();
502
503        match (key_builder, value_builder) {
504            (Ok(key_builder), Ok(value_builder)) => Ok(unsafe {
505                PrimitiveDictionaryBuilder::new_from_builders(key_builder, value_builder)
506            }),
507            (Err(key_array), Ok(mut value_builder)) => {
508                Err(Self::try_new(key_array, Arc::new(value_builder.finish())).unwrap())
509            }
510            (Ok(mut key_builder), Err(value_array)) => {
511                Err(Self::try_new(key_builder.finish(), Arc::new(value_array)).unwrap())
512            }
513            (Err(key_array), Err(value_array)) => {
514                Err(Self::try_new(key_array, Arc::new(value_array)).unwrap())
515            }
516        }
517    }
518
519    /// Applies an unary and infallible function to a mutable dictionary array.
520    /// Mutable dictionary array means that the buffers are not shared with other arrays.
521    /// As a result, this mutates the buffers directly without allocating new buffers.
522    ///
523    /// # Implementation
524    ///
525    /// This will apply the function for all dictionary values, including those on null slots.
526    /// This implies that the operation must be infallible for any value of the corresponding type
527    /// or this function may panic.
528    /// # Example
529    /// ```
530    /// # use std::sync::Arc;
531    /// # use arrow_array::{Array, ArrayAccessor, DictionaryArray, StringArray, types::{Int8Type, Int32Type}};
532    /// # use arrow_array::{Int8Array, Int32Array};
533    /// let values = Int32Array::from(vec![Some(10), Some(20), None]);
534    /// let keys = Int8Array::from_iter_values([0, 0, 1, 2]);
535    /// let dictionary = DictionaryArray::<Int8Type>::try_new(keys, Arc::new(values)).unwrap();
536    /// let c = dictionary.unary_mut::<_, Int32Type>(|x| x + 1).unwrap();
537    /// let typed = c.downcast_dict::<Int32Array>().unwrap();
538    /// assert_eq!(typed.value(0), 11);
539    /// assert_eq!(typed.value(1), 11);
540    /// assert_eq!(typed.value(2), 21);
541    /// ```
542    #[expect(clippy::result_large_err)]
543    pub fn unary_mut<F, V>(self, op: F) -> Result<DictionaryArray<K>, DictionaryArray<K>>
544    where
545        V: ArrowPrimitiveType,
546        F: Fn(V::Native) -> V::Native,
547    {
548        let mut builder: PrimitiveDictionaryBuilder<K, V> = self.into_primitive_dict_builder()?;
549        builder
550            .values_slice_mut()
551            .iter_mut()
552            .for_each(|v| *v = op(*v));
553        Ok(builder.finish())
554    }
555
556    /// Computes an occupancy mask for this dictionary's values
557    ///
558    /// For each value in [`Self::values`] the corresponding bit will be set in the
559    /// returned mask if it is referenced by a key in this [`DictionaryArray`]
560    pub fn occupancy(&self) -> BooleanBuffer {
561        let len = self.values.len();
562        let mut builder = BooleanBufferBuilder::new(len);
563        builder.resize(len);
564        let slice = builder.as_slice_mut();
565        match self.keys.nulls().filter(|n| n.null_count() > 0) {
566            Some(n) => {
567                let v = self.keys.values();
568                n.valid_indices()
569                    .for_each(|idx| set_bit(slice, v[idx].as_usize()))
570            }
571            None => {
572                let v = self.keys.values();
573                v.iter().for_each(|v| set_bit(slice, v.as_usize()))
574            }
575        }
576        builder.finish()
577    }
578}
579
580/// Constructs a `DictionaryArray` from an `ArrayData`
581impl<T: ArrowDictionaryKeyType> From<ArrayData> for DictionaryArray<T> {
582    fn from(data: ArrayData) -> Self {
583        let (data_type, len, nulls, offset, mut buffers, mut child_data) = data.into_parts();
584
585        assert_eq!(
586            buffers.len(),
587            1,
588            "DictionaryArray data should contain a single buffer only (keys)."
589        );
590        let buffer = buffers.pop().expect("checked above");
591        assert_eq!(
592            child_data.len(),
593            1,
594            "DictionaryArray should contain a single child array (values)."
595        );
596        let cd = child_data.pop().expect("checked above");
597
598        if let DataType::Dictionary(key_data_type, _) = &data_type {
599            assert_eq!(
600                &T::DATA_TYPE,
601                key_data_type.as_ref(),
602                "DictionaryArray's data type must match, expected {} got {}",
603                T::DATA_TYPE,
604                key_data_type
605            );
606
607            let values = make_array(cd);
608
609            // create a zero-copy of the keys' data
610            let keys = PrimitiveArray::<T>::new(ScalarBuffer::new(buffer, offset, len), nulls);
611
612            Self {
613                data_type,
614                keys,
615                values,
616                is_ordered: false,
617            }
618        } else {
619            panic!("DictionaryArray must have Dictionary data type.")
620        }
621    }
622}
623
624impl<T: ArrowDictionaryKeyType> From<DictionaryArray<T>> for ArrayData {
625    fn from(array: DictionaryArray<T>) -> Self {
626        let builder = array
627            .keys
628            .into_data()
629            .into_builder()
630            .data_type(array.data_type)
631            .child_data(vec![array.values.to_data()]);
632
633        unsafe { builder.build_unchecked() }
634    }
635}
636
637/// Constructs a `DictionaryArray` from an iterator of optional strings.
638///
639/// # Example:
640/// ```
641/// use arrow_array::{DictionaryArray, PrimitiveArray, StringArray, types::Int8Type};
642///
643/// let test = vec!["a", "a", "b", "c"];
644/// let array: DictionaryArray<Int8Type> = test
645///     .iter()
646///     .map(|&x| if x == "b" { None } else { Some(x) })
647///     .collect();
648/// assert_eq!(
649///     "DictionaryArray {keys: PrimitiveArray<Int8>\n[\n  0,\n  0,\n  null,\n  1,\n] values: StringArray\n[\n  \"a\",\n  \"c\",\n]}\n",
650///     format!("{:?}", array)
651/// );
652/// ```
653impl<'a, T: ArrowDictionaryKeyType> FromIterator<Option<&'a str>> for DictionaryArray<T> {
654    fn from_iter<I: IntoIterator<Item = Option<&'a str>>>(iter: I) -> Self {
655        let it = iter.into_iter();
656        let (lower, _) = it.size_hint();
657        let mut builder = StringDictionaryBuilder::with_capacity(lower, 256, 1024);
658        builder.extend(it);
659        builder.finish()
660    }
661}
662
663/// Constructs a `DictionaryArray` from an iterator of strings.
664///
665/// # Example:
666///
667/// ```
668/// use arrow_array::{DictionaryArray, PrimitiveArray, StringArray, types::Int8Type};
669///
670/// let test = vec!["a", "a", "b", "c"];
671/// let array: DictionaryArray<Int8Type> = test.into_iter().collect();
672/// assert_eq!(
673///     "DictionaryArray {keys: PrimitiveArray<Int8>\n[\n  0,\n  0,\n  1,\n  2,\n] values: StringArray\n[\n  \"a\",\n  \"b\",\n  \"c\",\n]}\n",
674///     format!("{:?}", array)
675/// );
676/// ```
677impl<'a, T: ArrowDictionaryKeyType> FromIterator<&'a str> for DictionaryArray<T> {
678    fn from_iter<I: IntoIterator<Item = &'a str>>(iter: I) -> Self {
679        let it = iter.into_iter();
680        let (lower, _) = it.size_hint();
681        let mut builder = StringDictionaryBuilder::with_capacity(lower, 256, 1024);
682        it.for_each(|i| {
683            builder
684                .append(i)
685                .expect("Unable to append a value to a dictionary array.");
686        });
687
688        builder.finish()
689    }
690}
691
692/// SAFETY: Correctly implements the contract of Arrow Arrays
693unsafe impl<T: ArrowDictionaryKeyType> Array for DictionaryArray<T> {
694    fn as_any(&self) -> &dyn Any {
695        self
696    }
697
698    fn to_data(&self) -> ArrayData {
699        self.clone().into()
700    }
701
702    fn into_data(self) -> ArrayData {
703        self.into()
704    }
705
706    fn data_type(&self) -> &DataType {
707        &self.data_type
708    }
709
710    fn slice(&self, offset: usize, length: usize) -> ArrayRef {
711        Arc::new(self.slice(offset, length))
712    }
713
714    fn len(&self) -> usize {
715        self.keys.len()
716    }
717
718    fn is_empty(&self) -> bool {
719        self.keys.is_empty()
720    }
721
722    fn shrink_to_fit(&mut self) {
723        self.keys.shrink_to_fit();
724        self.values.shrink_to_fit();
725    }
726
727    fn offset(&self) -> usize {
728        self.keys.offset()
729    }
730
731    fn nulls(&self) -> Option<&NullBuffer> {
732        self.keys.nulls()
733    }
734
735    fn logical_nulls(&self) -> Option<NullBuffer> {
736        match self.values.logical_nulls() {
737            None => self.nulls().cloned(),
738            Some(value_nulls) => {
739                let mut builder = BooleanBufferBuilder::new(self.len());
740                match self.keys.nulls() {
741                    Some(n) => builder.append_buffer(n.inner()),
742                    None => builder.append_n(self.len(), true),
743                }
744                for (idx, k) in self.keys.values().iter().enumerate() {
745                    let k = k.as_usize();
746                    // Check range to allow for nulls
747                    if k < value_nulls.len() && value_nulls.is_null(k) {
748                        builder.set_bit(idx, false);
749                    }
750                }
751                Some(builder.finish().into())
752            }
753        }
754    }
755
756    fn logical_null_count(&self) -> usize {
757        match (self.keys.nulls(), self.values.logical_nulls()) {
758            (None, None) => 0,
759            (Some(key_nulls), None) => key_nulls.null_count(),
760            (None, Some(value_nulls)) => self
761                .keys
762                .values()
763                .iter()
764                .filter(|k| value_nulls.is_null(k.as_usize()))
765                .count(),
766            (Some(key_nulls), Some(value_nulls)) => self
767                .keys
768                .values()
769                .iter()
770                .enumerate()
771                .filter(|(idx, k)| key_nulls.is_null(*idx) || value_nulls.is_null(k.as_usize()))
772                .count(),
773        }
774    }
775
776    fn is_nullable(&self) -> bool {
777        !self.is_empty() && (self.nulls().is_some() || self.values.is_nullable())
778    }
779
780    fn get_buffer_memory_size(&self) -> usize {
781        self.keys.get_buffer_memory_size() + self.values.get_buffer_memory_size()
782    }
783
784    fn get_array_memory_size(&self) -> usize {
785        std::mem::size_of::<Self>()
786            + self.keys.get_buffer_memory_size()
787            + self.values.get_array_memory_size()
788    }
789
790    #[cfg(feature = "pool")]
791    fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
792        self.keys.claim(pool);
793        self.values.claim(pool);
794    }
795}
796
797impl<T: ArrowDictionaryKeyType> std::fmt::Debug for DictionaryArray<T> {
798    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
799        writeln!(
800            f,
801            "DictionaryArray {{keys: {:?} values: {:?}}}",
802            self.keys, self.values
803        )
804    }
805}
806
807/// A [`DictionaryArray`] typed on its child values array
808///
809/// Implements [`ArrayAccessor`] allowing fast access to its elements
810///
811/// ```
812/// use arrow_array::{DictionaryArray, StringArray, types::Int32Type};
813///
814/// let orig = ["a", "b", "a", "b"];
815/// let dictionary = DictionaryArray::<Int32Type>::from_iter(orig);
816///
817/// // `TypedDictionaryArray` allows you to access the values directly
818/// let typed = dictionary.downcast_dict::<StringArray>().unwrap();
819///
820/// for (maybe_val, orig) in typed.into_iter().zip(orig) {
821///     assert_eq!(maybe_val.unwrap(), orig)
822/// }
823/// ```
824pub struct TypedDictionaryArray<'a, K: ArrowDictionaryKeyType, V> {
825    /// The dictionary array
826    dictionary: &'a DictionaryArray<K>,
827    /// The values of the dictionary
828    values: &'a V,
829}
830
831// Manually implement `Clone` to avoid `V: Clone` type constraint
832impl<K: ArrowDictionaryKeyType, V> Clone for TypedDictionaryArray<'_, K, V> {
833    fn clone(&self) -> Self {
834        *self
835    }
836}
837
838impl<K: ArrowDictionaryKeyType, V> Copy for TypedDictionaryArray<'_, K, V> {}
839
840impl<K: ArrowDictionaryKeyType, V> std::fmt::Debug for TypedDictionaryArray<'_, K, V> {
841    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
842        writeln!(f, "TypedDictionaryArray({:?})", self.dictionary)
843    }
844}
845
846impl<'a, K: ArrowDictionaryKeyType, V> TypedDictionaryArray<'a, K, V> {
847    /// Returns the keys of this [`TypedDictionaryArray`]
848    pub fn keys(&self) -> &'a PrimitiveArray<K> {
849        self.dictionary.keys()
850    }
851
852    /// Returns the values of this [`TypedDictionaryArray`]
853    pub fn values(&self) -> &'a V {
854        self.values
855    }
856}
857
858unsafe impl<K: ArrowDictionaryKeyType, V: Sync> Array for TypedDictionaryArray<'_, K, V> {
859    fn as_any(&self) -> &dyn Any {
860        self.dictionary
861    }
862
863    fn to_data(&self) -> ArrayData {
864        self.dictionary.to_data()
865    }
866
867    fn into_data(self) -> ArrayData {
868        self.dictionary.into_data()
869    }
870
871    fn data_type(&self) -> &DataType {
872        self.dictionary.data_type()
873    }
874
875    fn slice(&self, offset: usize, length: usize) -> ArrayRef {
876        Arc::new(self.dictionary.slice(offset, length))
877    }
878
879    fn len(&self) -> usize {
880        self.dictionary.len()
881    }
882
883    fn is_empty(&self) -> bool {
884        self.dictionary.is_empty()
885    }
886
887    fn offset(&self) -> usize {
888        self.dictionary.offset()
889    }
890
891    fn nulls(&self) -> Option<&NullBuffer> {
892        self.dictionary.nulls()
893    }
894
895    fn logical_nulls(&self) -> Option<NullBuffer> {
896        self.dictionary.logical_nulls()
897    }
898
899    fn logical_null_count(&self) -> usize {
900        self.dictionary.logical_null_count()
901    }
902
903    fn is_nullable(&self) -> bool {
904        self.dictionary.is_nullable()
905    }
906
907    fn get_buffer_memory_size(&self) -> usize {
908        self.dictionary.get_buffer_memory_size()
909    }
910
911    fn get_array_memory_size(&self) -> usize {
912        self.dictionary.get_array_memory_size()
913    }
914
915    #[cfg(feature = "pool")]
916    fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
917        self.dictionary.claim(pool);
918    }
919}
920
921impl<K, V> IntoIterator for TypedDictionaryArray<'_, K, V>
922where
923    K: ArrowDictionaryKeyType,
924    Self: ArrayAccessor,
925{
926    type Item = Option<<Self as ArrayAccessor>::Item>;
927    type IntoIter = ArrayIter<Self>;
928
929    fn into_iter(self) -> Self::IntoIter {
930        ArrayIter::new(self)
931    }
932}
933
934impl<'a, K, V> ArrayAccessor for TypedDictionaryArray<'a, K, V>
935where
936    K: ArrowDictionaryKeyType,
937    V: Sync + Send,
938    &'a V: ArrayAccessor,
939    <&'a V as ArrayAccessor>::Item: Default,
940{
941    type Item = <&'a V as ArrayAccessor>::Item;
942
943    fn value(&self, index: usize) -> Self::Item {
944        assert!(
945            index < self.len(),
946            "Trying to access an element at index {} from a TypedDictionaryArray of length {}",
947            index,
948            self.len()
949        );
950        unsafe { self.value_unchecked(index) }
951    }
952
953    unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
954        let val = unsafe { self.dictionary.keys.value_unchecked(index) };
955        let value_idx = val.as_usize();
956
957        // As dictionary keys are only verified for non-null indexes
958        // we must check the value is within bounds
959        match value_idx < self.values.len() {
960            true => unsafe { self.values.value_unchecked(value_idx) },
961            false => Default::default(),
962        }
963    }
964}
965
966/// A [`DictionaryArray`] with the key type erased
967///
968/// This can be used to efficiently implement kernels for all possible dictionary
969/// keys without needing to create specialized implementations for each key type
970///
971/// For example
972///
973/// ```
974/// # use arrow_array::*;
975/// # use arrow_array::cast::AsArray;
976/// # use arrow_array::builder::PrimitiveDictionaryBuilder;
977/// # use arrow_array::types::*;
978/// # use arrow_schema::ArrowError;
979/// # use std::sync::Arc;
980///
981/// fn to_string(a: &dyn Array) -> Result<ArrayRef, ArrowError> {
982///     if let Some(d) = a.as_any_dictionary_opt() {
983///         // Recursively handle dictionary input
984///         let r = to_string(d.values().as_ref())?;
985///         return Ok(d.with_values(r));
986///     }
987///     downcast_primitive_array! {
988///         a => Ok(Arc::new(a.iter().map(|x| x.map(|x| format!("{x:?}"))).collect::<StringArray>())),
989///         d => Err(ArrowError::InvalidArgumentError(format!("{d:?} not supported")))
990///     }
991/// }
992///
993/// let result = to_string(&Int32Array::from(vec![1, 2, 3])).unwrap();
994/// let actual = result.as_string::<i32>().iter().map(Option::unwrap).collect::<Vec<_>>();
995/// assert_eq!(actual, &["1", "2", "3"]);
996///
997/// let mut dict = PrimitiveDictionaryBuilder::<Int32Type, UInt16Type>::new();
998/// dict.extend([Some(1), Some(1), Some(2), Some(3), Some(2)]);
999/// let dict = dict.finish();
1000///
1001/// let r = to_string(&dict).unwrap();
1002/// let r = r.as_dictionary::<Int32Type>().downcast_dict::<StringArray>().unwrap();
1003/// assert_eq!(r.keys(), dict.keys()); // Keys are the same
1004///
1005/// let actual = r.into_iter().map(Option::unwrap).collect::<Vec<_>>();
1006/// assert_eq!(actual, &["1", "1", "2", "3", "2"]);
1007/// ```
1008///
1009/// See [`AsArray::as_any_dictionary_opt`] and [`AsArray::as_any_dictionary`]
1010pub trait AnyDictionaryArray: Array {
1011    /// Returns the primitive keys of this dictionary as an [`Array`]
1012    fn keys(&self) -> &dyn Array;
1013
1014    /// Returns the values of this dictionary
1015    fn values(&self) -> &ArrayRef;
1016
1017    /// Returns the keys of this dictionary as usize
1018    ///
1019    /// The values for nulls will be arbitrary, but are guaranteed
1020    /// to be in the range `0..self.values.len()`
1021    ///
1022    /// # Panics
1023    ///
1024    /// Panics if `values.len() == 0`
1025    fn normalized_keys(&self) -> Vec<usize>;
1026
1027    /// Create a new [`DictionaryArray`] replacing `values` with the new values
1028    ///
1029    /// See [`DictionaryArray::with_values`]
1030    fn with_values(&self, values: ArrayRef) -> ArrayRef;
1031}
1032
1033impl<K: ArrowDictionaryKeyType> AnyDictionaryArray for DictionaryArray<K> {
1034    fn keys(&self) -> &dyn Array {
1035        &self.keys
1036    }
1037
1038    fn values(&self) -> &ArrayRef {
1039        self.values()
1040    }
1041
1042    fn normalized_keys(&self) -> Vec<usize> {
1043        let v_len = self.values().len();
1044        assert_ne!(v_len, 0);
1045        let iter = self.keys().values().iter();
1046        iter.map(|x| x.as_usize().min(v_len - 1)).collect()
1047    }
1048
1049    fn with_values(&self, values: ArrayRef) -> ArrayRef {
1050        Arc::new(self.with_values(values))
1051    }
1052}
1053
1054#[cfg(test)]
1055mod tests {
1056    use super::*;
1057    use crate::cast::as_dictionary_array;
1058    use crate::{Decimal128Array, Int8Array, Int16Array, Int32Array, RunArray, UInt8Array};
1059    use arrow_buffer::{Buffer, ToByteSlice};
1060
1061    #[test]
1062    fn test_dictionary_array() {
1063        // Construct a value array
1064        let value_data = ArrayData::builder(DataType::Int8)
1065            .len(8)
1066            .add_buffer(Buffer::from(
1067                [10_i8, 11, 12, 13, 14, 15, 16, 17].to_byte_slice(),
1068            ))
1069            .build()
1070            .unwrap();
1071
1072        // Construct a buffer for value offsets, for the nested array:
1073        let keys = Buffer::from([2_i16, 3, 4].to_byte_slice());
1074
1075        // Construct a dictionary array from the above two
1076        let key_type = DataType::Int16;
1077        let value_type = DataType::Int8;
1078        let dict_data_type = DataType::Dictionary(Box::new(key_type), Box::new(value_type));
1079        let dict_data = ArrayData::builder(dict_data_type.clone())
1080            .len(3)
1081            .add_buffer(keys.clone())
1082            .add_child_data(value_data.clone())
1083            .build()
1084            .unwrap();
1085        let dict_array = Int16DictionaryArray::from(dict_data);
1086
1087        let values = dict_array.values();
1088        assert_eq!(value_data, values.to_data());
1089        assert_eq!(DataType::Int8, dict_array.value_type());
1090        assert_eq!(3, dict_array.len());
1091
1092        // Null count only makes sense in terms of the component arrays.
1093        assert_eq!(0, dict_array.null_count());
1094        assert_eq!(0, dict_array.values().null_count());
1095        assert_eq!(dict_array.keys(), &Int16Array::from(vec![2_i16, 3, 4]));
1096
1097        // Now test with a non-zero offset
1098        let dict_data = ArrayData::builder(dict_data_type)
1099            .len(2)
1100            .offset(1)
1101            .add_buffer(keys)
1102            .add_child_data(value_data.clone())
1103            .build()
1104            .unwrap();
1105        let dict_array = Int16DictionaryArray::from(dict_data);
1106
1107        let values = dict_array.values();
1108        assert_eq!(value_data, values.to_data());
1109        assert_eq!(DataType::Int8, dict_array.value_type());
1110        assert_eq!(2, dict_array.len());
1111        assert_eq!(dict_array.keys(), &Int16Array::from(vec![3_i16, 4]));
1112    }
1113
1114    #[test]
1115    fn test_dictionary_builder_append_many() {
1116        let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::new();
1117
1118        builder.append(1).unwrap();
1119        builder.append_n(2, 2).unwrap();
1120        builder.append_options(None, 2);
1121        builder.append_options(Some(3), 3);
1122
1123        let array = builder.finish();
1124
1125        let values = array
1126            .values()
1127            .as_primitive::<UInt32Type>()
1128            .iter()
1129            .map(Option::unwrap)
1130            .collect::<Vec<_>>();
1131        assert_eq!(values, &[1, 2, 3]);
1132        let keys = array.keys().iter().collect::<Vec<_>>();
1133        assert_eq!(
1134            keys,
1135            &[
1136                Some(0),
1137                Some(1),
1138                Some(1),
1139                None,
1140                None,
1141                Some(2),
1142                Some(2),
1143                Some(2)
1144            ]
1145        );
1146    }
1147
1148    #[test]
1149    fn test_string_dictionary_builder_append_many() {
1150        let mut builder = StringDictionaryBuilder::<Int8Type>::new();
1151
1152        builder.append("a").unwrap();
1153        builder.append_n("b", 2).unwrap();
1154        builder.append_options(None::<&str>, 2);
1155        builder.append_options(Some("c"), 3);
1156
1157        let array = builder.finish();
1158
1159        let values = array
1160            .values()
1161            .as_string::<i32>()
1162            .iter()
1163            .map(Option::unwrap)
1164            .collect::<Vec<_>>();
1165        assert_eq!(values, &["a", "b", "c"]);
1166        let keys = array.keys().iter().collect::<Vec<_>>();
1167        assert_eq!(
1168            keys,
1169            &[
1170                Some(0),
1171                Some(1),
1172                Some(1),
1173                None,
1174                None,
1175                Some(2),
1176                Some(2),
1177                Some(2)
1178            ]
1179        );
1180    }
1181
1182    #[test]
1183    fn test_dictionary_array_fmt_debug() {
1184        let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::with_capacity(3, 2);
1185        builder.append(12345678).unwrap();
1186        builder.append_null();
1187        builder.append(22345678).unwrap();
1188        let array = builder.finish();
1189        assert_eq!(
1190            "DictionaryArray {keys: PrimitiveArray<UInt8>\n[\n  0,\n  null,\n  1,\n] values: PrimitiveArray<UInt32>\n[\n  12345678,\n  22345678,\n]}\n",
1191            format!("{array:?}")
1192        );
1193
1194        let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::with_capacity(20, 2);
1195        for _ in 0..20 {
1196            builder.append(1).unwrap();
1197        }
1198        let array = builder.finish();
1199        assert_eq!(
1200            "DictionaryArray {keys: PrimitiveArray<UInt8>\n[\n  0,\n  0,\n  0,\n  0,\n  0,\n  0,\n  0,\n  0,\n  0,\n  0,\n  0,\n  0,\n  0,\n  0,\n  0,\n  0,\n  0,\n  0,\n  0,\n  0,\n] values: PrimitiveArray<UInt32>\n[\n  1,\n]}\n",
1201            format!("{array:?}")
1202        );
1203    }
1204
1205    #[test]
1206    fn test_dictionary_array_from_iter() {
1207        let test = vec!["a", "a", "b", "c"];
1208        let array: DictionaryArray<Int8Type> = test
1209            .iter()
1210            .map(|&x| if x == "b" { None } else { Some(x) })
1211            .collect();
1212        assert_eq!(
1213            "DictionaryArray {keys: PrimitiveArray<Int8>\n[\n  0,\n  0,\n  null,\n  1,\n] values: StringArray\n[\n  \"a\",\n  \"c\",\n]}\n",
1214            format!("{array:?}")
1215        );
1216
1217        let array: DictionaryArray<Int8Type> = test.into_iter().collect();
1218        assert_eq!(
1219            "DictionaryArray {keys: PrimitiveArray<Int8>\n[\n  0,\n  0,\n  1,\n  2,\n] values: StringArray\n[\n  \"a\",\n  \"b\",\n  \"c\",\n]}\n",
1220            format!("{array:?}")
1221        );
1222    }
1223
1224    #[test]
1225    fn test_dictionary_array_reverse_lookup_key() {
1226        let test = vec!["a", "a", "b", "c"];
1227        let array: DictionaryArray<Int8Type> = test.into_iter().collect();
1228
1229        assert_eq!(array.lookup_key("c"), Some(2));
1230
1231        // Direction of building a dictionary is the iterator direction
1232        let test = vec!["t3", "t3", "t2", "t2", "t1", "t3", "t4", "t1", "t0"];
1233        let array: DictionaryArray<Int8Type> = test.into_iter().collect();
1234
1235        assert_eq!(array.lookup_key("t1"), Some(2));
1236        assert_eq!(array.lookup_key("non-existent"), None);
1237    }
1238
1239    #[test]
1240    fn test_dictionary_keys_as_primitive_array() {
1241        let test = vec!["a", "b", "c", "a"];
1242        let array: DictionaryArray<Int8Type> = test.into_iter().collect();
1243
1244        let keys = array.keys();
1245        assert_eq!(&DataType::Int8, keys.data_type());
1246        assert_eq!(0, keys.null_count());
1247        assert_eq!(&[0, 1, 2, 0], keys.values());
1248    }
1249
1250    #[test]
1251    fn test_dictionary_keys_as_primitive_array_with_null() {
1252        let test = vec![Some("a"), None, Some("b"), None, None, Some("a")];
1253        let array: DictionaryArray<Int32Type> = test.into_iter().collect();
1254
1255        let keys = array.keys();
1256        assert_eq!(&DataType::Int32, keys.data_type());
1257        assert_eq!(3, keys.null_count());
1258
1259        assert!(keys.is_valid(0));
1260        assert!(!keys.is_valid(1));
1261        assert!(keys.is_valid(2));
1262        assert!(!keys.is_valid(3));
1263        assert!(!keys.is_valid(4));
1264        assert!(keys.is_valid(5));
1265
1266        assert_eq!(0, keys.value(0));
1267        assert_eq!(1, keys.value(2));
1268        assert_eq!(0, keys.value(5));
1269    }
1270
1271    #[test]
1272    fn test_dictionary_all_nulls() {
1273        let test = vec![None, None, None];
1274        let array: DictionaryArray<Int32Type> = test.into_iter().collect();
1275        array
1276            .into_data()
1277            .validate_full()
1278            .expect("All null array has valid array data");
1279    }
1280
1281    #[test]
1282    fn test_dictionary_iter() {
1283        // Construct a value array
1284        let values = Int8Array::from_iter_values([10_i8, 11, 12, 13, 14, 15, 16, 17]);
1285        let keys = Int16Array::from_iter_values([2_i16, 3, 4]);
1286
1287        // Construct a dictionary array from the above two
1288        let dict_array = DictionaryArray::new(keys, Arc::new(values));
1289
1290        let mut key_iter = dict_array.keys_iter();
1291        assert_eq!(2, key_iter.next().unwrap().unwrap());
1292        assert_eq!(3, key_iter.next().unwrap().unwrap());
1293        assert_eq!(4, key_iter.next().unwrap().unwrap());
1294        assert!(key_iter.next().is_none());
1295
1296        let mut iter = dict_array
1297            .values()
1298            .as_any()
1299            .downcast_ref::<Int8Array>()
1300            .unwrap()
1301            .take_iter(dict_array.keys_iter());
1302
1303        assert_eq!(12, iter.next().unwrap().unwrap());
1304        assert_eq!(13, iter.next().unwrap().unwrap());
1305        assert_eq!(14, iter.next().unwrap().unwrap());
1306        assert!(iter.next().is_none());
1307    }
1308
1309    #[test]
1310    fn test_dictionary_iter_with_null() {
1311        let test = vec![Some("a"), None, Some("b"), None, None, Some("a")];
1312        let array: DictionaryArray<Int32Type> = test.into_iter().collect();
1313
1314        let mut iter = array
1315            .values()
1316            .as_any()
1317            .downcast_ref::<StringArray>()
1318            .unwrap()
1319            .take_iter(array.keys_iter());
1320
1321        assert_eq!("a", iter.next().unwrap().unwrap());
1322        assert!(iter.next().unwrap().is_none());
1323        assert_eq!("b", iter.next().unwrap().unwrap());
1324        assert!(iter.next().unwrap().is_none());
1325        assert!(iter.next().unwrap().is_none());
1326        assert_eq!("a", iter.next().unwrap().unwrap());
1327        assert!(iter.next().is_none());
1328    }
1329
1330    #[test]
1331    fn test_dictionary_key() {
1332        let keys = Int8Array::from(vec![Some(2), None, Some(1)]);
1333        let values = StringArray::from(vec!["foo", "bar", "baz", "blarg"]);
1334
1335        let array = DictionaryArray::new(keys, Arc::new(values));
1336        assert_eq!(array.key(0), Some(2));
1337        assert_eq!(array.key(1), None);
1338        assert_eq!(array.key(2), Some(1));
1339    }
1340
1341    #[test]
1342    fn test_try_new() {
1343        let values: StringArray = [Some("foo"), Some("bar"), Some("baz")]
1344            .into_iter()
1345            .collect();
1346        let keys: Int32Array = [Some(0), Some(2), None, Some(1)].into_iter().collect();
1347
1348        let array = DictionaryArray::new(keys, Arc::new(values));
1349        assert_eq!(array.keys().data_type(), &DataType::Int32);
1350        assert_eq!(array.values().data_type(), &DataType::Utf8);
1351
1352        assert_eq!(array.null_count(), 1);
1353        assert_eq!(array.logical_null_count(), 1);
1354
1355        assert!(array.keys().is_valid(0));
1356        assert!(array.keys().is_valid(1));
1357        assert!(array.keys().is_null(2));
1358        assert!(array.keys().is_valid(3));
1359
1360        assert_eq!(array.keys().value(0), 0);
1361        assert_eq!(array.keys().value(1), 2);
1362        assert_eq!(array.keys().value(3), 1);
1363
1364        assert_eq!(
1365            "DictionaryArray {keys: PrimitiveArray<Int32>\n[\n  0,\n  2,\n  null,\n  1,\n] values: StringArray\n[\n  \"foo\",\n  \"bar\",\n  \"baz\",\n]}\n",
1366            format!("{array:?}")
1367        );
1368    }
1369
1370    #[test]
1371    #[should_panic(expected = "Invalid dictionary key 3 at index 1, expected 0 <= key < 2")]
1372    fn test_try_new_index_too_large() {
1373        let values: StringArray = [Some("foo"), Some("bar")].into_iter().collect();
1374        // dictionary only has 2 values, so offset 3 is out of bounds
1375        let keys: Int32Array = [Some(0), Some(3)].into_iter().collect();
1376        DictionaryArray::new(keys, Arc::new(values));
1377    }
1378
1379    #[test]
1380    #[should_panic(expected = "Invalid dictionary key -100 at index 0, expected 0 <= key < 2")]
1381    fn test_try_new_index_too_small() {
1382        let values: StringArray = [Some("foo"), Some("bar")].into_iter().collect();
1383        let keys: Int32Array = std::iter::once(Some(-100)).collect();
1384        DictionaryArray::new(keys, Arc::new(values));
1385    }
1386
1387    #[test]
1388    #[should_panic(expected = "DictionaryArray's data type must match, expected Int64 got Int32")]
1389    fn test_from_array_data_validation() {
1390        let a = DictionaryArray::<Int32Type>::from_iter(["32"]);
1391        let _ = DictionaryArray::<Int64Type>::from(a.into_data());
1392    }
1393
1394    #[test]
1395    fn test_into_primitive_dict_builder() {
1396        let values = Int32Array::from_iter_values([10_i32, 12, 15]);
1397        let keys = Int8Array::from_iter_values([1_i8, 0, 2, 0]);
1398
1399        let dict_array = DictionaryArray::new(keys, Arc::new(values));
1400
1401        let boxed: ArrayRef = Arc::new(dict_array);
1402        let col: DictionaryArray<Int8Type> = as_dictionary_array(&boxed).clone();
1403
1404        drop(boxed);
1405
1406        let mut builder = col.into_primitive_dict_builder::<Int32Type>().unwrap();
1407
1408        let slice = builder.values_slice_mut();
1409        assert_eq!(slice, &[10, 12, 15]);
1410
1411        slice[0] = 4;
1412        slice[1] = 2;
1413        slice[2] = 1;
1414
1415        let values = Int32Array::from_iter_values([4_i32, 2, 1]);
1416        let keys = Int8Array::from_iter_values([1_i8, 0, 2, 0]);
1417
1418        let expected = DictionaryArray::new(keys, Arc::new(values));
1419
1420        let new_array = builder.finish();
1421        assert_eq!(expected, new_array);
1422    }
1423
1424    #[test]
1425    fn test_into_primitive_dict_builder_cloned_array() {
1426        let values = Int32Array::from_iter_values([10_i32, 12, 15]);
1427        let keys = Int8Array::from_iter_values([1_i8, 0, 2, 0]);
1428
1429        let dict_array = DictionaryArray::new(keys, Arc::new(values));
1430
1431        let boxed: ArrayRef = Arc::new(dict_array);
1432
1433        let col: DictionaryArray<Int8Type> = DictionaryArray::<Int8Type>::from(boxed.to_data());
1434        let err = col.into_primitive_dict_builder::<Int32Type>();
1435
1436        let returned = err.unwrap_err();
1437
1438        let values = Int32Array::from_iter_values([10_i32, 12, 15]);
1439        let keys = Int8Array::from_iter_values([1_i8, 0, 2, 0]);
1440
1441        let expected = DictionaryArray::new(keys, Arc::new(values));
1442        assert_eq!(expected, returned);
1443    }
1444
1445    #[test]
1446    fn test_unary_mut_keeps_value_type() {
1447        let values = || {
1448            Decimal128Array::from(vec![123, 456])
1449                .with_precision_and_scale(5, 2)
1450                .unwrap()
1451        };
1452        let keys = Int8Array::from_iter_values([0, 1, 0]);
1453        let dict = DictionaryArray::new(keys, Arc::new(values()));
1454        let data_type = dict.data_type().clone();
1455        let updated = dict.unary_mut::<_, Decimal128Type>(|x| x + 1).unwrap();
1456        assert_eq!(updated.data_type(), &data_type);
1457
1458        // When the keys are shared, the original array is returned
1459        let keys = Int8Array::from_iter_values([0, 1, 0]);
1460        let dict = DictionaryArray::new(keys.clone(), Arc::new(values()));
1461        let returned = dict.unary_mut::<_, Decimal128Type>(|x| x + 1).unwrap_err();
1462        assert_eq!(returned, DictionaryArray::new(keys, Arc::new(values())));
1463    }
1464
1465    #[test]
1466    fn test_occupancy() {
1467        let keys = Int32Array::new((100..200).collect(), None);
1468        let values = Int32Array::from(vec![0; 1024]);
1469        let dict = DictionaryArray::new(keys, Arc::new(values));
1470        for (idx, v) in dict.occupancy().iter().enumerate() {
1471            let expected = (100..200).contains(&idx);
1472            assert_eq!(v, expected, "{idx}");
1473        }
1474
1475        let keys = Int32Array::new(
1476            (0..100).collect(),
1477            Some((0..100).map(|x| x % 4 == 0).collect()),
1478        );
1479        let values = Int32Array::from(vec![0; 1024]);
1480        let dict = DictionaryArray::new(keys, Arc::new(values));
1481        for (idx, v) in dict.occupancy().iter().enumerate() {
1482            let expected = idx % 4 == 0 && idx < 100;
1483            assert_eq!(v, expected, "{idx}");
1484        }
1485    }
1486
1487    #[test]
1488    fn test_iterator_nulls() {
1489        let keys = Int32Array::new(
1490            vec![0, 700, 1, 2].into(),
1491            Some(NullBuffer::from(vec![true, false, true, true])),
1492        );
1493        let values = Int32Array::from(vec![Some(50), None, Some(2)]);
1494        let dict = DictionaryArray::new(keys, Arc::new(values));
1495        let values: Vec<_> = dict
1496            .downcast_dict::<Int32Array>()
1497            .unwrap()
1498            .into_iter()
1499            .collect();
1500        assert_eq!(values, &[Some(50), None, None, Some(2)])
1501    }
1502
1503    #[test]
1504    fn test_logical_nulls() -> Result<(), ArrowError> {
1505        let values = Arc::new(RunArray::try_new(
1506            &Int32Array::from(vec![1, 3, 7]),
1507            &Int32Array::from(vec![Some(1), None, Some(3)]),
1508        )?) as ArrayRef;
1509
1510        // For this test to be meaningful, the values array need to have different nulls and logical nulls
1511        assert_eq!(values.null_count(), 0);
1512        assert_eq!(values.logical_null_count(), 2);
1513
1514        // Construct a trivial dictionary with 1-1 mapping to underlying array
1515        let dictionary = DictionaryArray::<Int8Type>::try_new(
1516            Int8Array::from((0..values.len()).map(|i| i as i8).collect::<Vec<_>>()),
1517            Arc::clone(&values),
1518        )?;
1519
1520        // No keys are null
1521        assert_eq!(dictionary.null_count(), 0);
1522        // Dictionary array values are logically nullable
1523        assert_eq!(dictionary.logical_null_count(), values.logical_null_count());
1524        assert_eq!(dictionary.logical_nulls(), values.logical_nulls());
1525        assert!(dictionary.is_nullable());
1526
1527        // Construct a trivial dictionary with 1-1 mapping to underlying array except that key 0 is nulled out
1528        let dictionary = DictionaryArray::<Int8Type>::try_new(
1529            Int8Array::from(
1530                (0..values.len())
1531                    .map(|i| i as i8)
1532                    .map(|i| if i == 0 { None } else { Some(i) })
1533                    .collect::<Vec<_>>(),
1534            ),
1535            Arc::clone(&values),
1536        )?;
1537
1538        // One key is null
1539        assert_eq!(dictionary.null_count(), 1);
1540
1541        // Dictionary array values are logically nullable
1542        assert_eq!(
1543            dictionary.logical_null_count(),
1544            values.logical_null_count() + 1
1545        );
1546        assert!(dictionary.is_nullable());
1547
1548        Ok(())
1549    }
1550
1551    #[test]
1552    fn test_normalized_keys() {
1553        let values = vec![132, 0, 1].into();
1554        let nulls = NullBuffer::from(vec![false, true, true]);
1555        let keys = Int32Array::new(values, Some(nulls));
1556        let dictionary = DictionaryArray::new(keys, Arc::new(Int32Array::new_null(2)));
1557        assert_eq!(&dictionary.normalized_keys(), &[1, 0, 1])
1558    }
1559
1560    #[test]
1561    fn test_all_null_dict() {
1562        let all_null_dict_arr = DictionaryArray::try_new(
1563            UInt8Array::new_null(10),
1564            Arc::new(StringArray::from_iter_values(["a"])),
1565        );
1566        assert!(all_null_dict_arr.is_ok())
1567    }
1568}
arrow_array/array/dictionary_array.rs

arrow_array/array/
dictionary_array.rs
