polars_arrow/array/
iterator.rs

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

polars_arrow/array/
iterator.rs
