Skip to main content

polars_core/chunked_array/iterator/
mod.rs

1use arrow::array::*;
2
3use crate::prelude::*;
4
5pub mod par;
6
7impl<T> ChunkedArray<T>
8where
9    T: PolarsDataType,
10{
11    #[inline]
12    pub fn iter(&self) -> impl PolarsIterator<Item = Option<T::Physical<'_>>> {
13        // SAFETY: we set the correct length of the iterator.
14        unsafe {
15            self.downcast_iter()
16                .flat_map(|arr| arr.iter())
17                .trust_my_length(self.len())
18        }
19    }
20
21    #[inline]
22    pub fn no_null_iter(&self) -> impl PolarsIterator<Item = T::Physical<'_>> {
23        // SAFETY: we set the correct length of the iterator.
24        unsafe {
25            self.downcast_iter()
26                .flat_map(|arr| arr.values_iter())
27                .trust_my_length(self.len())
28        }
29    }
30}
31
32/// A [`PolarsIterator`] is an iterator over a [`ChunkedArray`] which contains polars types. A [`PolarsIterator`]
33/// must implement [`ExactSizeIterator`] and [`DoubleEndedIterator`].
34pub trait PolarsIterator:
35    ExactSizeIterator + DoubleEndedIterator + Send + Sync + TrustedLen
36{
37}
38
39/// Implement [`PolarsIterator`] for every iterator that implements the needed traits.
40impl<T: ?Sized> PolarsIterator for T where
41    T: ExactSizeIterator + DoubleEndedIterator + Send + Sync + TrustedLen
42{
43}
44
45impl ListChunked {
46    pub fn series_iter(&self) -> impl PolarsIterator<Item = Option<Series>> {
47        let dtype = self.inner_dtype();
48        unsafe {
49            self.downcast_iter()
50                .flat_map(|arr| arr.iter())
51                .trust_my_length(self.len())
52                .map(move |arr| {
53                    arr.map(|arr| {
54                        Series::from_chunks_and_dtype_unchecked(PlSmallStr::EMPTY, vec![arr], dtype)
55                    })
56                })
57        }
58    }
59
60    pub fn no_null_series_iter(&self) -> impl PolarsIterator<Item = Series> {
61        let inner_type = self.inner_dtype();
62        unsafe {
63            self.downcast_iter()
64                .flat_map(|arr| arr.values_iter())
65                .map(move |arr| {
66                    Series::from_chunks_and_dtype_unchecked(
67                        PlSmallStr::EMPTY,
68                        vec![arr],
69                        inner_type,
70                    )
71                })
72                .trust_my_length(self.len())
73        }
74    }
75}
76
77#[cfg(feature = "dtype-array")]
78impl ArrayChunked {
79    pub fn series_iter(&self) -> impl PolarsIterator<Item = Option<Series>> {
80        let dtype = self.inner_dtype();
81        unsafe {
82            self.downcast_iter()
83                .flat_map(|arr| arr.iter())
84                .trust_my_length(self.len())
85                .map(move |arr| {
86                    arr.map(|arr| {
87                        Series::from_chunks_and_dtype_unchecked(PlSmallStr::EMPTY, vec![arr], dtype)
88                    })
89                })
90        }
91    }
92}