polars_arrow/array/dictionary/
mod.rs

1use std::hash::Hash;
2use std::hint::unreachable_unchecked;
3
4use crate::bitmap::Bitmap;
5use crate::bitmap::utils::{BitmapIter, ZipValidity};
6use crate::datatypes::{ArrowDataType, IntegerType};
7use crate::scalar::{Scalar, new_scalar};
8use crate::trusted_len::TrustedLen;
9use crate::types::NativeType;
10
11mod ffi;
12pub(super) mod fmt;
13mod iterator;
14mod mutable;
15use crate::array::specification::check_indexes_unchecked;
16mod typed_iterator;
17mod value_map;
18
19pub use iterator::*;
20pub use mutable::*;
21use polars_error::{PolarsResult, polars_bail};
22
23use super::primitive::PrimitiveArray;
24use super::specification::check_indexes;
25use super::{Array, Splitable, new_empty_array, new_null_array};
26use crate::array::dictionary::typed_iterator::{
27    DictValue, DictionaryIterTyped, DictionaryValuesIterTyped,
28};
29
30/// Trait denoting [`NativeType`]s that can be used as keys of a dictionary.
31/// # Safety
32///
33/// Any implementation of this trait must ensure that `always_fits_usize` only
34/// returns `true` if all values succeeds on `value::try_into::<usize>().unwrap()`.
35pub unsafe trait DictionaryKey: NativeType + TryInto<usize> + TryFrom<usize> + Hash {
36    /// The corresponding [`IntegerType`] of this key
37    const KEY_TYPE: IntegerType;
38    const MAX_USIZE_VALUE: usize;
39
40    /// Represents this key as a `usize`.
41    ///
42    /// # Safety
43    /// The caller _must_ have checked that the value can be cast to `usize`.
44    #[inline]
45    unsafe fn as_usize(self) -> usize {
46        match self.try_into() {
47            Ok(v) => v,
48            Err(_) => unreachable_unchecked(),
49        }
50    }
51
52    /// Create a key from a `usize` without checking bounds.
53    ///
54    /// # Safety
55    /// The caller _must_ have checked that the value can be created from a `usize`.
56    #[inline]
57    unsafe fn from_usize_unchecked(x: usize) -> Self {
58        debug_assert!(Self::try_from(x).is_ok());
59        unsafe { Self::try_from(x).unwrap_unchecked() }
60    }
61
62    /// If the key type always can be converted to `usize`.
63    fn always_fits_usize() -> bool {
64        false
65    }
66}
67
68unsafe impl DictionaryKey for i8 {
69    const KEY_TYPE: IntegerType = IntegerType::Int8;
70    const MAX_USIZE_VALUE: usize = i8::MAX as usize;
71}
72unsafe impl DictionaryKey for i16 {
73    const KEY_TYPE: IntegerType = IntegerType::Int16;
74    const MAX_USIZE_VALUE: usize = i16::MAX as usize;
75}
76unsafe impl DictionaryKey for i32 {
77    const KEY_TYPE: IntegerType = IntegerType::Int32;
78    const MAX_USIZE_VALUE: usize = i32::MAX as usize;
79}
80unsafe impl DictionaryKey for i64 {
81    const KEY_TYPE: IntegerType = IntegerType::Int64;
82    const MAX_USIZE_VALUE: usize = i64::MAX as usize;
83}
84unsafe impl DictionaryKey for i128 {
85    const KEY_TYPE: IntegerType = IntegerType::Int128;
86    const MAX_USIZE_VALUE: usize = i128::MAX as usize;
87}
88unsafe impl DictionaryKey for u8 {
89    const KEY_TYPE: IntegerType = IntegerType::UInt8;
90    const MAX_USIZE_VALUE: usize = u8::MAX as usize;
91
92    fn always_fits_usize() -> bool {
93        true
94    }
95}
96unsafe impl DictionaryKey for u16 {
97    const KEY_TYPE: IntegerType = IntegerType::UInt16;
98    const MAX_USIZE_VALUE: usize = u16::MAX as usize;
99
100    fn always_fits_usize() -> bool {
101        true
102    }
103}
104unsafe impl DictionaryKey for u32 {
105    const KEY_TYPE: IntegerType = IntegerType::UInt32;
106    const MAX_USIZE_VALUE: usize = u32::MAX as usize;
107
108    fn always_fits_usize() -> bool {
109        true
110    }
111}
112unsafe impl DictionaryKey for u64 {
113    const KEY_TYPE: IntegerType = IntegerType::UInt64;
114    const MAX_USIZE_VALUE: usize = u64::MAX as usize;
115
116    #[cfg(target_pointer_width = "64")]
117    fn always_fits_usize() -> bool {
118        true
119    }
120}
121unsafe impl DictionaryKey for u128 {
122    const KEY_TYPE: IntegerType = IntegerType::UInt128;
123    const MAX_USIZE_VALUE: usize = u128::MAX as usize;
124}
125
126/// An [`Array`] whose values are stored as indices. This [`Array`] is useful when the cardinality of
127/// values is low compared to the length of the [`Array`].
128///
129/// # Safety
130/// This struct guarantees that each item of [`DictionaryArray::keys`] is castable to `usize` and
131/// its value is smaller than [`DictionaryArray::values`]`.len()`. In other words, you can safely
132/// use `unchecked` calls to retrieve the values
133#[derive(Clone)]
134pub struct DictionaryArray<K: DictionaryKey> {
135    dtype: ArrowDataType,
136    keys: PrimitiveArray<K>,
137    values: Box<dyn Array>,
138}
139
140fn check_dtype(
141    key_type: IntegerType,
142    dtype: &ArrowDataType,
143    values_dtype: &ArrowDataType,
144) -> PolarsResult<()> {
145    if let ArrowDataType::Dictionary(key, value, _) = dtype.to_storage() {
146        if *key != key_type {
147            polars_bail!(ComputeError: "DictionaryArray must be initialized with a DataType::Dictionary whose integer is compatible to its keys")
148        }
149        if value.as_ref().to_storage() != values_dtype.to_storage() {
150            polars_bail!(ComputeError: "DictionaryArray must be initialized with a DataType::Dictionary whose value is equal to its values")
151        }
152    } else {
153        polars_bail!(ComputeError: "DictionaryArray must be initialized with logical DataType::Dictionary")
154    }
155    Ok(())
156}
157
158impl<K: DictionaryKey> DictionaryArray<K> {
159    /// Returns a new [`DictionaryArray`].
160    /// # Implementation
161    /// This function is `O(N)` where `N` is the length of keys
162    /// # Errors
163    /// This function errors iff
164    /// * the `dtype`'s logical type is not a `DictionaryArray`
165    /// * the `dtype`'s keys is not compatible with `keys`
166    /// * the `dtype`'s values's dtype is not equal with `values.dtype()`
167    /// * any of the keys's values is not represented in `usize` or is `>= values.len()`
168    pub fn try_new(
169        dtype: ArrowDataType,
170        keys: PrimitiveArray<K>,
171        values: Box<dyn Array>,
172    ) -> PolarsResult<Self> {
173        check_dtype(K::KEY_TYPE, &dtype, values.dtype())?;
174
175        if keys.null_count() != keys.len() {
176            if K::always_fits_usize() {
177                // SAFETY: we just checked that conversion to `usize` always
178                // succeeds
179                unsafe { check_indexes_unchecked(keys.values(), values.len()) }?;
180            } else {
181                check_indexes(keys.values(), values.len())?;
182            }
183        }
184
185        Ok(Self {
186            dtype,
187            keys,
188            values,
189        })
190    }
191
192    /// Returns a new [`DictionaryArray`].
193    /// # Implementation
194    /// This function is `O(N)` where `N` is the length of keys
195    /// # Errors
196    /// This function errors iff
197    /// * any of the keys's values is not represented in `usize` or is `>= values.len()`
198    pub fn try_from_keys(
199        keys: PrimitiveArray<K>,
200        values: Box<dyn Array>,
201        ordered: bool,
202    ) -> PolarsResult<Self> {
203        let dtype = Self::default_dtype(values.dtype().clone(), ordered);
204        Self::try_new(dtype, keys, values)
205    }
206
207    /// Returns a new [`DictionaryArray`].
208    /// # Errors
209    /// This function errors iff
210    /// * the `dtype`'s logical type is not a `DictionaryArray`
211    /// * the `dtype`'s keys is not compatible with `keys`
212    /// * the `dtype`'s values's dtype is not equal with `values.dtype()`
213    ///
214    /// # Safety
215    /// The caller must ensure that every keys's values is represented in `usize` and is `< values.len()`
216    pub unsafe fn try_new_unchecked(
217        dtype: ArrowDataType,
218        keys: PrimitiveArray<K>,
219        values: Box<dyn Array>,
220    ) -> PolarsResult<Self> {
221        check_dtype(K::KEY_TYPE, &dtype, values.dtype())?;
222
223        Ok(Self {
224            dtype,
225            keys,
226            values,
227        })
228    }
229
230    /// Returns a new empty [`DictionaryArray`].
231    pub fn new_empty(dtype: ArrowDataType) -> Self {
232        let values = Self::try_get_child(&dtype).unwrap();
233        let values = new_empty_array(values.clone());
234        Self::try_new(
235            dtype,
236            PrimitiveArray::<K>::new_empty(K::PRIMITIVE.into()),
237            values,
238        )
239        .unwrap()
240    }
241
242    /// Returns an [`DictionaryArray`] whose all elements are null
243    #[inline]
244    pub fn new_null(dtype: ArrowDataType, length: usize) -> Self {
245        let values = Self::try_get_child(&dtype).unwrap();
246        let values = new_null_array(values.clone(), 1);
247        Self::try_new(
248            dtype,
249            PrimitiveArray::<K>::new_null(K::PRIMITIVE.into(), length),
250            values,
251        )
252        .unwrap()
253    }
254
255    /// Returns an iterator of [`Option<Box<dyn Scalar>>`].
256    /// # Implementation
257    /// This function will allocate a new [`Scalar`] per item and is usually not performant.
258    /// Consider calling `keys_iter` and `values`, downcasting `values`, and iterating over that.
259    pub fn iter(
260        &self,
261    ) -> ZipValidity<Box<dyn Scalar>, DictionaryValuesIter<'_, K>, BitmapIter<'_>> {
262        ZipValidity::new_with_validity(DictionaryValuesIter::new(self), self.keys.validity())
263    }
264
265    /// Returns an iterator of [`Box<dyn Scalar>`]
266    /// # Implementation
267    /// This function will allocate a new [`Scalar`] per item and is usually not performant.
268    /// Consider calling `keys_iter` and `values`, downcasting `values`, and iterating over that.
269    pub fn values_iter(&self) -> DictionaryValuesIter<'_, K> {
270        DictionaryValuesIter::new(self)
271    }
272
273    /// Returns an iterator over the values [`V::IterValue`].
274    ///
275    /// # Panics
276    ///
277    /// Panics if the keys of this [`DictionaryArray`] has any nulls.
278    /// If they do [`DictionaryArray::iter_typed`] should be used.
279    pub fn values_iter_typed<V: DictValue>(
280        &self,
281    ) -> PolarsResult<DictionaryValuesIterTyped<'_, K, V>> {
282        let keys = &self.keys;
283        assert_eq!(keys.null_count(), 0);
284        let values = self.values.as_ref();
285        let values = V::downcast_values(values)?;
286        Ok(DictionaryValuesIterTyped::new(keys, values))
287    }
288
289    /// Returns an iterator over the optional values of  [`Option<V::IterValue>`].
290    pub fn iter_typed<V: DictValue>(&self) -> PolarsResult<DictionaryIterTyped<'_, K, V>> {
291        let keys = &self.keys;
292        let values = self.values.as_ref();
293        let values = V::downcast_values(values)?;
294        Ok(DictionaryIterTyped::new(keys, values))
295    }
296
297    /// Returns the [`ArrowDataType`] of this [`DictionaryArray`]
298    #[inline]
299    pub fn dtype(&self) -> &ArrowDataType {
300        &self.dtype
301    }
302
303    /// Returns whether the values of this [`DictionaryArray`] are ordered
304    #[inline]
305    pub fn is_ordered(&self) -> bool {
306        match self.dtype.to_storage() {
307            ArrowDataType::Dictionary(_, _, is_ordered) => *is_ordered,
308            _ => unreachable!(),
309        }
310    }
311
312    pub(crate) fn default_dtype(values_datatype: ArrowDataType, ordered: bool) -> ArrowDataType {
313        ArrowDataType::Dictionary(K::KEY_TYPE, Box::new(values_datatype), ordered)
314    }
315
316    /// Slices this [`DictionaryArray`].
317    /// # Panics
318    /// iff `offset + length > self.len()`.
319    pub fn slice(&mut self, offset: usize, length: usize) {
320        self.keys.slice(offset, length);
321    }
322
323    /// Slices this [`DictionaryArray`].
324    ///
325    /// # Safety
326    /// Safe iff `offset + length <= self.len()`.
327    pub unsafe fn slice_unchecked(&mut self, offset: usize, length: usize) {
328        self.keys.slice_unchecked(offset, length);
329    }
330
331    impl_sliced!();
332
333    /// Returns this [`DictionaryArray`] with a new validity.
334    /// # Panic
335    /// This function panics iff `validity.len() != self.len()`.
336    #[must_use]
337    pub fn with_validity(mut self, validity: Option<Bitmap>) -> Self {
338        self.set_validity(validity);
339        self
340    }
341
342    /// Sets the validity of the keys of this [`DictionaryArray`].
343    /// # Panics
344    /// This function panics iff `validity.len() != self.len()`.
345    pub fn set_validity(&mut self, validity: Option<Bitmap>) {
346        self.keys.set_validity(validity);
347    }
348
349    impl_into_array!();
350
351    /// Returns the length of this array
352    #[inline]
353    pub fn len(&self) -> usize {
354        self.keys.len()
355    }
356
357    /// The optional validity. Equivalent to `self.keys().validity()`.
358    #[inline]
359    pub fn validity(&self) -> Option<&Bitmap> {
360        self.keys.validity()
361    }
362
363    /// Returns the keys of the [`DictionaryArray`]. These keys can be used to fetch values
364    /// from `values`.
365    #[inline]
366    pub fn keys(&self) -> &PrimitiveArray<K> {
367        &self.keys
368    }
369
370    /// Returns an iterator of the keys' values of the [`DictionaryArray`] as `usize`
371    #[inline]
372    pub fn keys_values_iter(&self) -> impl TrustedLen<Item = usize> + Clone + '_ {
373        // SAFETY: invariant of the struct
374        self.keys.values_iter().map(|x| unsafe { x.as_usize() })
375    }
376
377    /// Returns an iterator of the keys' of the [`DictionaryArray`] as `usize`
378    #[inline]
379    pub fn keys_iter(&self) -> impl TrustedLen<Item = Option<usize>> + Clone + '_ {
380        // SAFETY: invariant of the struct
381        self.keys.iter().map(|x| x.map(|x| unsafe { x.as_usize() }))
382    }
383
384    /// Returns the keys' value of the [`DictionaryArray`] as `usize`
385    /// # Panics
386    /// This function panics iff `index >= self.len()`
387    #[inline]
388    pub fn key_value(&self, index: usize) -> usize {
389        // SAFETY: invariant of the struct
390        unsafe { self.keys.values()[index].as_usize() }
391    }
392
393    /// Returns the values of the [`DictionaryArray`].
394    #[inline]
395    pub fn values(&self) -> &Box<dyn Array> {
396        &self.values
397    }
398
399    /// Returns the value of the [`DictionaryArray`] at position `i`.
400    /// # Implementation
401    /// This function will allocate a new [`Scalar`] and is usually not performant.
402    /// Consider calling `keys` and `values`, downcasting `values`, and iterating over that.
403    /// # Panic
404    /// This function panics iff `index >= self.len()`
405    #[inline]
406    pub fn value(&self, index: usize) -> Box<dyn Scalar> {
407        // SAFETY: invariant of this struct
408        let index = unsafe { self.keys.value(index).as_usize() };
409        new_scalar(self.values.as_ref(), index)
410    }
411
412    pub(crate) fn try_get_child(dtype: &ArrowDataType) -> PolarsResult<&ArrowDataType> {
413        Ok(match dtype.to_storage() {
414            ArrowDataType::Dictionary(_, values, _) => values.as_ref(),
415            _ => {
416                polars_bail!(ComputeError: "Dictionaries must be initialized with DataType::Dictionary")
417            },
418        })
419    }
420
421    pub fn take(self) -> (ArrowDataType, PrimitiveArray<K>, Box<dyn Array>) {
422        (self.dtype, self.keys, self.values)
423    }
424}
425
426impl<K: DictionaryKey> Array for DictionaryArray<K> {
427    impl_common_array!();
428
429    fn validity(&self) -> Option<&Bitmap> {
430        self.keys.validity()
431    }
432
433    #[inline]
434    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
435        Box::new(self.clone().with_validity(validity))
436    }
437}
438
439impl<K: DictionaryKey> Splitable for DictionaryArray<K> {
440    fn check_bound(&self, offset: usize) -> bool {
441        offset < self.len()
442    }
443
444    unsafe fn _split_at_unchecked(&self, offset: usize) -> (Self, Self) {
445        let (lhs_keys, rhs_keys) = unsafe { Splitable::split_at_unchecked(&self.keys, offset) };
446
447        (
448            Self {
449                dtype: self.dtype.clone(),
450                keys: lhs_keys,
451                values: self.values.clone(),
452            },
453            Self {
454                dtype: self.dtype.clone(),
455                keys: rhs_keys,
456                values: self.values.clone(),
457            },
458        )
459    }
460}
polars_arrow/array/dictionary/mod.rs

polars_arrow/array/dictionary/
mod.rs
