polars_core/chunked_array/object/
iterator.rs

Help
1use arrow::array::Array;
2use arrow::trusted_len::TrustedLen;
3
4use crate::chunked_array::object::{ObjectArray, PolarsObject};
5
6/// An iterator that returns Some(T) or None, that can be used on any ObjectArray
7// Note: This implementation is based on std's [Vec]s' [IntoIter].
8pub struct ObjectIter<'a, T: PolarsObject> {
9    array: &'a ObjectArray<T>,
10    current: usize,
11    current_end: usize,
12}
13
14impl<'a, T: PolarsObject> ObjectIter<'a, T> {
15    /// create a new iterator
16    pub fn new(array: &'a ObjectArray<T>) -> Self {
17        ObjectIter::<T> {
18            array,
19            current: 0,
20            current_end: array.len(),
21        }
22    }
23}
24
25impl<'a, T: PolarsObject> std::iter::Iterator for ObjectIter<'a, T> {
26    type Item = Option<&'a T>;
27
28    #[inline]
29    fn next(&mut self) -> Option<Self::Item> {
30        if self.current == self.current_end {
31            None
32        // SAFETY:
33        // Se comment below
34        } else if unsafe { self.array.is_null_unchecked(self.current) } {
35            self.current += 1;
36            Some(None)
37        } else {
38            let old = self.current;
39            self.current += 1;
40            // SAFETY:
41            // we just checked bounds in `self.current_end == self.current`
42            // this is safe on the premise that this struct is initialized with
43            // current = array.len()
44            // and that current_end is ever only decremented
45            unsafe { Some(Some(self.array.value_unchecked(old))) }
46        }
47    }
48
49    fn size_hint(&self) -> (usize, Option<usize>) {
50        (
51            self.array.len() - self.current,
52            Some(self.array.len() - self.current),
53        )
54    }
55}
56
57impl<T: PolarsObject> std::iter::DoubleEndedIterator for ObjectIter<'_, T> {
58    fn next_back(&mut self) -> Option<Self::Item> {
59        if self.current_end == self.current {
60            None
61        } else {
62            self.current_end -= 1;
63            Some(if self.array.is_null(self.current_end) {
64                None
65            } else {
66                // SAFETY:
67                // we just checked bounds in `self.current_end == self.current`
68                // this is safe on the premise that this struct is initialized with
69                // current = array.len()
70                // and that current_end is ever only decremented
71                unsafe { Some(self.array.value_unchecked(self.current_end)) }
72            })
73        }
74    }
75}
76
77/// all arrays have known size.
78impl<T: PolarsObject> std::iter::ExactSizeIterator for ObjectIter<'_, T> {}
79
80impl<'a, T: PolarsObject> IntoIterator for &'a ObjectArray<T> {
81    type Item = Option<&'a T>;
82    type IntoIter = ObjectIter<'a, T>;
83
84    fn into_iter(self) -> Self::IntoIter {
85        ObjectIter::<'a, T>::new(self)
86    }
87}
88
89pub struct OwnedObjectIter<T: PolarsObject> {
90    array: ObjectArray<T>,
91    current: usize,
92    current_end: usize,
93}
94
95impl<T: PolarsObject> OwnedObjectIter<T> {
96    /// create a new iterator
97    pub fn new(array: ObjectArray<T>) -> Self {
98        let current_end = array.len();
99        OwnedObjectIter::<T> {
100            array,
101            current: 0,
102            current_end,
103        }
104    }
105}
106
107unsafe impl<T: PolarsObject> TrustedLen for OwnedObjectIter<T> {}
108
109impl<T: PolarsObject> ObjectArray<T> {
110    pub(crate) fn into_iter_cloned(self) -> OwnedObjectIter<T> {
111        OwnedObjectIter::<T>::new(self)
112    }
113}
114impl<T: PolarsObject> std::iter::Iterator for OwnedObjectIter<T> {
115    type Item = Option<T>;
116
117    #[inline]
118    fn next(&mut self) -> Option<Self::Item> {
119        if self.current == self.current_end {
120            None
121        // SAFETY:
122        // Se comment below
123        } else if unsafe { self.array.is_null_unchecked(self.current) } {
124            self.current += 1;
125            Some(None)
126        } else {
127            let old = self.current;
128            self.current += 1;
129            // SAFETY:
130            // we just checked bounds in `self.current_end == self.current`
131            // this is safe on the premise that this struct is initialized with
132            // current = array.len()
133            // and that current_end is ever only decremented
134            unsafe { Some(Some(self.array.value_unchecked(old).clone())) }
135        }
136    }
137
138    fn size_hint(&self) -> (usize, Option<usize>) {
139        (
140            self.array.len() - self.current,
141            Some(self.array.len() - self.current),
142        )
143    }
144}
polars_core/chunked_array/object/iterator.rs

polars_core/chunked_array/object/
iterator.rs
