polars_arrow/array/dictionary/
iterator.rs

1use super::{DictionaryArray, DictionaryKey};
2use crate::bitmap::utils::{BitmapIter, ZipValidity};
3use crate::scalar::Scalar;
4use crate::trusted_len::TrustedLen;
5
6/// Iterator of values of an `ListArray`.
7pub struct DictionaryValuesIter<'a, K: DictionaryKey> {
8    array: &'a DictionaryArray<K>,
9    index: usize,
10    end: usize,
11}
12
13impl<'a, K: DictionaryKey> DictionaryValuesIter<'a, K> {
14    #[inline]
15    pub fn new(array: &'a DictionaryArray<K>) -> Self {
16        Self {
17            array,
18            index: 0,
19            end: array.len(),
20        }
21    }
22}
23
24impl<K: DictionaryKey> Iterator for DictionaryValuesIter<'_, K> {
25    type Item = Box<dyn Scalar>;
26
27    #[inline]
28    fn next(&mut self) -> Option<Self::Item> {
29        if self.index == self.end {
30            return None;
31        }
32        let old = self.index;
33        self.index += 1;
34        Some(self.array.value(old))
35    }
36
37    #[inline]
38    fn size_hint(&self) -> (usize, Option<usize>) {
39        (self.end - self.index, Some(self.end - self.index))
40    }
41}
42
43unsafe impl<K: DictionaryKey> TrustedLen for DictionaryValuesIter<'_, K> {}
44
45impl<K: DictionaryKey> DoubleEndedIterator for DictionaryValuesIter<'_, K> {
46    #[inline]
47    fn next_back(&mut self) -> Option<Self::Item> {
48        if self.index == self.end {
49            None
50        } else {
51            self.end -= 1;
52            Some(self.array.value(self.end))
53        }
54    }
55}
56
57type ValuesIter<'a, K> = DictionaryValuesIter<'a, K>;
58type ZipIter<'a, K> = ZipValidity<Box<dyn Scalar>, ValuesIter<'a, K>, BitmapIter<'a>>;
59
60impl<'a, K: DictionaryKey> IntoIterator for &'a DictionaryArray<K> {
61    type Item = Option<Box<dyn Scalar>>;
62    type IntoIter = ZipIter<'a, K>;
63
64    fn into_iter(self) -> Self::IntoIter {
65        self.iter()
66    }
67}