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

arrow_array/array/
dictionary_array.rs
