polars_arrow/array/
iterator.rs

1use crate::bitmap::iterator::TrueIdxIter;
2use crate::bitmap::Bitmap;
3use crate::trusted_len::TrustedLen;
4
5mod private {
6    pub trait Sealed {}
7
8    impl<'a, T: super::ArrayAccessor<'a> + ?Sized> Sealed for T {}
9}
10
11/// Sealed trait representing access to a value of an array.
12/// # Safety
13/// Implementers of this trait guarantee that
14/// `value_unchecked` is safe when called up to `len`
15pub unsafe trait ArrayAccessor<'a>: private::Sealed {
16    type Item: 'a;
17    /// # Safety
18    /// The index must be in-bounds in the array.
19    unsafe fn value_unchecked(&'a self, index: usize) -> Self::Item;
20    fn len(&self) -> usize;
21}
22
23/// Iterator of values of an [`ArrayAccessor`].
24#[derive(Debug, Clone)]
25pub struct ArrayValuesIter<'a, A: ArrayAccessor<'a>> {
26    array: &'a A,
27    index: usize,
28    end: usize,
29}
30
31impl<'a, A: ArrayAccessor<'a>> ArrayValuesIter<'a, A> {
32    /// Creates a new [`ArrayValuesIter`]
33    #[inline]
34    pub fn new(array: &'a A) -> Self {
35        Self {
36            array,
37            index: 0,
38            end: array.len(),
39        }
40    }
41}
42
43impl<'a, A: ArrayAccessor<'a>> Iterator for ArrayValuesIter<'a, A> {
44    type Item = A::Item;
45
46    #[inline]
47    fn next(&mut self) -> Option<Self::Item> {
48        if self.index == self.end {
49            return None;
50        }
51        let old = self.index;
52        self.index += 1;
53        Some(unsafe { self.array.value_unchecked(old) })
54    }
55
56    #[inline]
57    fn size_hint(&self) -> (usize, Option<usize>) {
58        (self.end - self.index, Some(self.end - self.index))
59    }
60
61    #[inline]
62    fn nth(&mut self, n: usize) -> Option<Self::Item> {
63        let new_index = self.index + n;
64        if new_index > self.end {
65            self.index = self.end;
66            None
67        } else {
68            self.index = new_index;
69            self.next()
70        }
71    }
72}
73
74impl<'a, A: ArrayAccessor<'a>> DoubleEndedIterator for ArrayValuesIter<'a, A> {
75    #[inline]
76    fn next_back(&mut self) -> Option<Self::Item> {
77        if self.index == self.end {
78            None
79        } else {
80            self.end -= 1;
81            Some(unsafe { self.array.value_unchecked(self.end) })
82        }
83    }
84}
85
86unsafe impl<'a, A: ArrayAccessor<'a>> TrustedLen for ArrayValuesIter<'a, A> {}
87impl<'a, A: ArrayAccessor<'a>> ExactSizeIterator for ArrayValuesIter<'a, A> {}
88
89pub struct NonNullValuesIter<'a, A: ?Sized> {
90    accessor: &'a A,
91    idxs: TrueIdxIter<'a>,
92}
93
94impl<'a, A: ArrayAccessor<'a> + ?Sized> NonNullValuesIter<'a, A> {
95    pub fn new(accessor: &'a A, validity: Option<&'a Bitmap>) -> Self {
96        Self {
97            idxs: TrueIdxIter::new(accessor.len(), validity),
98            accessor,
99        }
100    }
101}
102
103impl<'a, A: ArrayAccessor<'a> + ?Sized> Iterator for NonNullValuesIter<'a, A> {
104    type Item = A::Item;
105
106    #[inline]
107    fn next(&mut self) -> Option<Self::Item> {
108        if let Some(i) = self.idxs.next() {
109            return Some(unsafe { self.accessor.value_unchecked(i) });
110        }
111        None
112    }
113
114    fn size_hint(&self) -> (usize, Option<usize>) {
115        self.idxs.size_hint()
116    }
117}
118
119unsafe impl<'a, A: ArrayAccessor<'a> + ?Sized> TrustedLen for NonNullValuesIter<'a, A> {}
120
121impl<A: ?Sized> Clone for NonNullValuesIter<'_, A> {
122    fn clone(&self) -> Self {
123        Self {
124            accessor: self.accessor,
125            idxs: self.idxs.clone(),
126        }
127    }
128}
polars_arrow/array/iterator.rs

polars_arrow/array/
iterator.rs
