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

arrow_array/array/
dictionary_array.rs
