polars_core/series/
iterator.rs

1use crate::prelude::any_value::arr_to_any_value;
2use crate::prelude::*;
3use crate::utils::NoNull;
4
5macro_rules! from_iterator {
6    ($native:ty, $variant:ident) => {
7        impl FromIterator<Option<$native>> for Series {
8            fn from_iter<I: IntoIterator<Item = Option<$native>>>(iter: I) -> Self {
9                let ca: ChunkedArray<$variant> = iter.into_iter().collect();
10                ca.into_series()
11            }
12        }
13
14        impl FromIterator<$native> for Series {
15            fn from_iter<I: IntoIterator<Item = $native>>(iter: I) -> Self {
16                let ca: NoNull<ChunkedArray<$variant>> = iter.into_iter().collect();
17                ca.into_inner().into_series()
18            }
19        }
20
21        impl<'a> FromIterator<&'a $native> for Series {
22            fn from_iter<I: IntoIterator<Item = &'a $native>>(iter: I) -> Self {
23                let ca: ChunkedArray<$variant> = iter.into_iter().map(|v| Some(*v)).collect();
24                ca.into_series()
25            }
26        }
27    };
28}
29
30#[cfg(feature = "dtype-u8")]
31from_iterator!(u8, UInt8Type);
32#[cfg(feature = "dtype-u16")]
33from_iterator!(u16, UInt16Type);
34from_iterator!(u32, UInt32Type);
35from_iterator!(u64, UInt64Type);
36#[cfg(feature = "dtype-i8")]
37from_iterator!(i8, Int8Type);
38#[cfg(feature = "dtype-i16")]
39from_iterator!(i16, Int16Type);
40from_iterator!(i32, Int32Type);
41from_iterator!(i64, Int64Type);
42from_iterator!(f32, Float32Type);
43from_iterator!(f64, Float64Type);
44from_iterator!(bool, BooleanType);
45
46impl<'a> FromIterator<Option<&'a str>> for Series {
47    fn from_iter<I: IntoIterator<Item = Option<&'a str>>>(iter: I) -> Self {
48        let ca: StringChunked = iter.into_iter().collect();
49        ca.into_series()
50    }
51}
52
53impl<'a> FromIterator<&'a str> for Series {
54    fn from_iter<I: IntoIterator<Item = &'a str>>(iter: I) -> Self {
55        let ca: StringChunked = iter.into_iter().collect();
56        ca.into_series()
57    }
58}
59
60impl FromIterator<Option<String>> for Series {
61    fn from_iter<T: IntoIterator<Item = Option<String>>>(iter: T) -> Self {
62        let ca: StringChunked = iter.into_iter().collect();
63        ca.into_series()
64    }
65}
66
67impl FromIterator<String> for Series {
68    fn from_iter<I: IntoIterator<Item = String>>(iter: I) -> Self {
69        let ca: StringChunked = iter.into_iter().collect();
70        ca.into_series()
71    }
72}
73
74impl Series {
75    /// Iterate over [`Series`] as [`AnyValue`].
76    ///
77    /// # Panics
78    /// This will panic if the array is not rechunked first.
79    pub fn iter(&self) -> SeriesIter<'_> {
80        let arrays = self.chunks();
81        SeriesIter {
82            idx_in_cur_arr: 0,
83            cur_arr_idx: 0,
84            cur_arr_len: arrays[0].len(),
85            arrays,
86            dtype: self.dtype(),
87            total_elems_in_remaining_arrays: self.len(),
88        }
89    }
90}
91
92pub struct SeriesIter<'a> {
93    arrays: &'a [Box<dyn Array>],
94    dtype: &'a DataType,
95    idx_in_cur_arr: usize,
96    cur_arr_len: usize,
97    cur_arr_idx: usize,
98    total_elems_in_remaining_arrays: usize,
99}
100
101impl<'a> Iterator for SeriesIter<'a> {
102    type Item = AnyValue<'a>;
103
104    #[inline]
105    fn next(&mut self) -> Option<Self::Item> {
106        loop {
107            if self.idx_in_cur_arr < self.cur_arr_len {
108                let arr = unsafe { self.arrays.get_unchecked(self.cur_arr_idx) };
109                let ret = unsafe { arr_to_any_value(&**arr, self.idx_in_cur_arr, self.dtype) };
110                self.idx_in_cur_arr += 1;
111                return Some(ret);
112            }
113
114            if self.cur_arr_idx + 1 < self.arrays.len() {
115                self.total_elems_in_remaining_arrays -= self.cur_arr_len;
116                self.cur_arr_idx += 1;
117                self.idx_in_cur_arr = 0;
118                let arr = unsafe { self.arrays.get_unchecked(self.cur_arr_idx) };
119                self.cur_arr_len = arr.len();
120            } else {
121                return None;
122            }
123        }
124    }
125
126    fn size_hint(&self) -> (usize, Option<usize>) {
127        let len = self.total_elems_in_remaining_arrays - self.idx_in_cur_arr;
128        (len, Some(len))
129    }
130}
131
132impl ExactSizeIterator for SeriesIter<'_> {}
133
134#[cfg(test)]
135mod test {
136    use crate::prelude::*;
137
138    #[test]
139    fn test_iter() {
140        let a = Series::new("age".into(), [23, 71, 9].as_ref());
141        let _b = a.i32().unwrap().iter().map(|opt_v| opt_v.map(|v| v * 2));
142    }
143
144    #[test]
145    fn test_iter_str() {
146        let data = [Some("John"), Some("Doe"), None];
147        let a: Series = data.into_iter().collect();
148        let b = Series::new("".into(), data);
149        assert_eq!(a, b);
150    }
151
152    #[test]
153    fn test_iter_string() {
154        let data = [Some("John".to_string()), Some("Doe".to_string()), None];
155        let a: Series = data.clone().into_iter().collect();
156        let b = Series::new("".into(), data);
157        assert_eq!(a, b);
158    }
159}
polars_core/series/iterator.rs

polars_core/series/
iterator.rs
