polars_arrow/bitmap/utils/
zip_validity.rs

1use crate::bitmap::utils::BitmapIter;
2use crate::bitmap::Bitmap;
3use crate::trusted_len::TrustedLen;
4
5/// An [`Iterator`] over validity and values.
6#[derive(Debug, Clone)]
7pub struct ZipValidityIter<T, I, V>
8where
9    I: Iterator<Item = T>,
10    V: Iterator<Item = bool>,
11{
12    values: I,
13    validity: V,
14}
15
16impl<T, I, V> ZipValidityIter<T, I, V>
17where
18    I: Iterator<Item = T>,
19    V: Iterator<Item = bool>,
20{
21    /// Creates a new [`ZipValidityIter`].
22    /// # Panics
23    /// This function panics if the size_hints of the iterators are different
24    pub fn new(values: I, validity: V) -> Self {
25        assert_eq!(values.size_hint(), validity.size_hint());
26        Self { values, validity }
27    }
28}
29
30impl<T, I, V> Iterator for ZipValidityIter<T, I, V>
31where
32    I: Iterator<Item = T>,
33    V: Iterator<Item = bool>,
34{
35    type Item = Option<T>;
36
37    #[inline]
38    fn next(&mut self) -> Option<Self::Item> {
39        let value = self.values.next();
40        let is_valid = self.validity.next();
41        is_valid
42            .zip(value)
43            .map(|(is_valid, value)| is_valid.then(|| value))
44    }
45
46    #[inline]
47    fn size_hint(&self) -> (usize, Option<usize>) {
48        self.values.size_hint()
49    }
50
51    #[inline]
52    fn nth(&mut self, n: usize) -> Option<Self::Item> {
53        let value = self.values.nth(n);
54        let is_valid = self.validity.nth(n);
55        is_valid
56            .zip(value)
57            .map(|(is_valid, value)| is_valid.then(|| value))
58    }
59}
60
61impl<T, I, V> DoubleEndedIterator for ZipValidityIter<T, I, V>
62where
63    I: DoubleEndedIterator<Item = T>,
64    V: DoubleEndedIterator<Item = bool>,
65{
66    #[inline]
67    fn next_back(&mut self) -> Option<Self::Item> {
68        let value = self.values.next_back();
69        let is_valid = self.validity.next_back();
70        is_valid
71            .zip(value)
72            .map(|(is_valid, value)| is_valid.then(|| value))
73    }
74}
75
76unsafe impl<T, I, V> TrustedLen for ZipValidityIter<T, I, V>
77where
78    I: TrustedLen<Item = T>,
79    V: TrustedLen<Item = bool>,
80{
81}
82
83impl<T, I, V> ExactSizeIterator for ZipValidityIter<T, I, V>
84where
85    I: ExactSizeIterator<Item = T>,
86    V: ExactSizeIterator<Item = bool>,
87{
88}
89
90/// An [`Iterator`] over [`Option<T>`]
91/// This enum can be used in two distinct ways:
92/// * as an iterator, via `Iterator::next`
93/// * as an enum of two iterators, via `match self`
94///
95/// The latter allows specializalizing to when there are no nulls
96#[derive(Debug, Clone)]
97pub enum ZipValidity<T, I, V>
98where
99    I: Iterator<Item = T>,
100    V: Iterator<Item = bool>,
101{
102    /// There are no null values
103    Required(I),
104    /// There are null values
105    Optional(ZipValidityIter<T, I, V>),
106}
107
108impl<T, I, V> ZipValidity<T, I, V>
109where
110    I: Iterator<Item = T>,
111    V: Iterator<Item = bool>,
112{
113    /// Returns a new [`ZipValidity`]
114    pub fn new(values: I, validity: Option<V>) -> Self {
115        match validity {
116            Some(validity) => Self::Optional(ZipValidityIter::new(values, validity)),
117            _ => Self::Required(values),
118        }
119    }
120}
121
122impl<'a, T, I> ZipValidity<T, I, BitmapIter<'a>>
123where
124    I: Iterator<Item = T>,
125{
126    /// Returns a new [`ZipValidity`] and drops the `validity` if all values
127    /// are valid.
128    pub fn new_with_validity(values: I, validity: Option<&'a Bitmap>) -> Self {
129        // only if the validity has nulls we take the optional branch.
130        match validity.and_then(|validity| (validity.unset_bits() > 0).then(|| validity.iter())) {
131            Some(validity) => Self::Optional(ZipValidityIter::new(values, validity)),
132            _ => Self::Required(values),
133        }
134    }
135}
136
137impl<T, I, V> Iterator for ZipValidity<T, I, V>
138where
139    I: Iterator<Item = T>,
140    V: Iterator<Item = bool>,
141{
142    type Item = Option<T>;
143
144    #[inline]
145    fn next(&mut self) -> Option<Self::Item> {
146        match self {
147            Self::Required(values) => values.next().map(Some),
148            Self::Optional(zipped) => zipped.next(),
149        }
150    }
151
152    #[inline]
153    fn size_hint(&self) -> (usize, Option<usize>) {
154        match self {
155            Self::Required(values) => values.size_hint(),
156            Self::Optional(zipped) => zipped.size_hint(),
157        }
158    }
159
160    #[inline]
161    fn nth(&mut self, n: usize) -> Option<Self::Item> {
162        match self {
163            Self::Required(values) => values.nth(n).map(Some),
164            Self::Optional(zipped) => zipped.nth(n),
165        }
166    }
167}
168
169impl<T, I, V> DoubleEndedIterator for ZipValidity<T, I, V>
170where
171    I: DoubleEndedIterator<Item = T>,
172    V: DoubleEndedIterator<Item = bool>,
173{
174    #[inline]
175    fn next_back(&mut self) -> Option<Self::Item> {
176        match self {
177            Self::Required(values) => values.next_back().map(Some),
178            Self::Optional(zipped) => zipped.next_back(),
179        }
180    }
181}
182
183impl<T, I, V> ExactSizeIterator for ZipValidity<T, I, V>
184where
185    I: ExactSizeIterator<Item = T>,
186    V: ExactSizeIterator<Item = bool>,
187{
188}
189
190unsafe impl<T, I, V> TrustedLen for ZipValidity<T, I, V>
191where
192    I: TrustedLen<Item = T>,
193    V: TrustedLen<Item = bool>,
194{
195}
196
197impl<T, I, V> ZipValidity<T, I, V>
198where
199    I: Iterator<Item = T>,
200    V: Iterator<Item = bool>,
201{
202    /// Unwrap into an iterator that has no null values.
203    pub fn unwrap_required(self) -> I {
204        match self {
205            ZipValidity::Required(i) => i,
206            _ => panic!("Could not 'unwrap_required'. 'ZipValidity' iterator has nulls."),
207        }
208    }
209
210    /// Unwrap into an iterator that has null values.
211    pub fn unwrap_optional(self) -> ZipValidityIter<T, I, V> {
212        match self {
213            ZipValidity::Optional(i) => i,
214            _ => panic!("Could not 'unwrap_optional'. 'ZipValidity' iterator has no nulls."),
215        }
216    }
217}
polars_arrow/bitmap/utils/zip_validity.rs

polars_arrow/bitmap/utils/
zip_validity.rs
