polars_arrow/array/struct_/
iterator.rs

1use super::StructArray;
2use crate::bitmap::utils::{BitmapIter, ZipValidity};
3use crate::scalar::{new_scalar, Scalar};
4use crate::trusted_len::TrustedLen;
5
6pub struct StructValueIter<'a> {
7    array: &'a StructArray,
8    index: usize,
9    end: usize,
10}
11
12impl<'a> StructValueIter<'a> {
13    #[inline]
14    pub fn new(array: &'a StructArray) -> Self {
15        Self {
16            array,
17            index: 0,
18            end: array.len(),
19        }
20    }
21}
22
23impl Iterator for StructValueIter<'_> {
24    type Item = Vec<Box<dyn Scalar>>;
25
26    #[inline]
27    fn next(&mut self) -> Option<Self::Item> {
28        if self.index == self.end {
29            return None;
30        }
31        let old = self.index;
32        self.index += 1;
33
34        // SAFETY:
35        // self.end is maximized by the length of the array
36        Some(
37            self.array
38                .values()
39                .iter()
40                .map(|v| new_scalar(v.as_ref(), old))
41                .collect(),
42        )
43    }
44
45    #[inline]
46    fn size_hint(&self) -> (usize, Option<usize>) {
47        (self.end - self.index, Some(self.end - self.index))
48    }
49}
50
51unsafe impl TrustedLen for StructValueIter<'_> {}
52
53impl DoubleEndedIterator for StructValueIter<'_> {
54    #[inline]
55    fn next_back(&mut self) -> Option<Self::Item> {
56        if self.index == self.end {
57            None
58        } else {
59            self.end -= 1;
60
61            // SAFETY:
62            // self.end is maximized by the length of the array
63            Some(
64                self.array
65                    .values()
66                    .iter()
67                    .map(|v| new_scalar(v.as_ref(), self.end))
68                    .collect(),
69            )
70        }
71    }
72}
73
74type ValuesIter<'a> = StructValueIter<'a>;
75type ZipIter<'a> = ZipValidity<Vec<Box<dyn Scalar>>, ValuesIter<'a>, BitmapIter<'a>>;
76
77impl<'a> IntoIterator for &'a StructArray {
78    type Item = Option<Vec<Box<dyn Scalar>>>;
79    type IntoIter = ZipIter<'a>;
80
81    fn into_iter(self) -> Self::IntoIter {
82        self.iter()
83    }
84}
85
86impl<'a> StructArray {
87    /// Returns an iterator of `Option<Box<dyn Array>>`
88    pub fn iter(&'a self) -> ZipIter<'a> {
89        ZipValidity::new_with_validity(StructValueIter::new(self), self.validity())
90    }
91
92    /// Returns an iterator of `Box<dyn Array>`
93    pub fn values_iter(&'a self) -> ValuesIter<'a> {
94        StructValueIter::new(self)
95    }
96}
polars_arrow/array/struct_/iterator.rs

polars_arrow/array/struct_/
iterator.rs
