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
74pub type SeriesPhysIter<'a> = Box<dyn ExactSizeIterator<Item = AnyValue<'a>> + 'a>;
75
76impl Series {
77    /// iterate over [`Series`] as [`AnyValue`].
78    ///
79    /// # Panics
80    /// This will panic if the array is not rechunked first.
81    pub fn iter(&self) -> SeriesIter<'_> {
82        let dtype = self.dtype();
83        #[cfg(feature = "object")]
84        assert!(
85            !matches!(dtype, DataType::Object(_)),
86            "object dtype not supported in Series.iter"
87        );
88        assert_eq!(self.chunks().len(), 1, "impl error");
89        let arr = &*self.chunks()[0];
90        let len = arr.len();
91        SeriesIter {
92            arr,
93            dtype,
94            idx: 0,
95            len,
96        }
97    }
98
99    pub fn phys_iter(&self) -> SeriesPhysIter<'_> {
100        let dtype = self.dtype();
101        let phys_dtype = dtype.to_physical();
102
103        assert_eq!(dtype, &phys_dtype, "impl error");
104        assert_eq!(self.chunks().len(), 1, "impl error");
105        #[cfg(feature = "object")]
106        assert!(
107            !matches!(dtype, DataType::Object(_)),
108            "object dtype not supported in Series.iter"
109        );
110        let arr = &*self.chunks()[0];
111
112        if phys_dtype.is_primitive_numeric() {
113            if arr.null_count() == 0 {
114                with_match_physical_numeric_type!(phys_dtype, |$T| {
115                        let arr = arr.as_any().downcast_ref::<PrimitiveArray<$T>>().unwrap();
116                        let values = arr.values().as_slice();
117                        Box::new(values.iter().map(|&value| AnyValue::from(value))) as Box<dyn ExactSizeIterator<Item=AnyValue<'_>> + '_>
118                })
119            } else {
120                with_match_physical_numeric_type!(phys_dtype, |$T| {
121                        let arr = arr.as_any().downcast_ref::<PrimitiveArray<$T>>().unwrap();
122                        Box::new(arr.iter().map(|value| {
123
124                        match value {
125                            Some(value) => AnyValue::from(*value),
126                            None => AnyValue::Null
127                        }
128
129                    })) as Box<dyn ExactSizeIterator<Item=AnyValue<'_>> + '_>
130                })
131            }
132        } else {
133            match dtype {
134                DataType::String => {
135                    let arr = arr.as_any().downcast_ref::<Utf8ViewArray>().unwrap();
136                    if arr.null_count() == 0 {
137                        Box::new(arr.values_iter().map(AnyValue::String))
138                            as Box<dyn ExactSizeIterator<Item = AnyValue<'_>> + '_>
139                    } else {
140                        let zipvalid = arr.iter();
141                        Box::new(zipvalid.unwrap_optional().map(|v| match v {
142                            Some(value) => AnyValue::String(value),
143                            None => AnyValue::Null,
144                        }))
145                            as Box<dyn ExactSizeIterator<Item = AnyValue<'_>> + '_>
146                    }
147                },
148                DataType::Boolean => {
149                    let arr = arr.as_any().downcast_ref::<BooleanArray>().unwrap();
150                    if arr.null_count() == 0 {
151                        Box::new(arr.values_iter().map(AnyValue::Boolean))
152                            as Box<dyn ExactSizeIterator<Item = AnyValue<'_>> + '_>
153                    } else {
154                        let zipvalid = arr.iter();
155                        Box::new(zipvalid.unwrap_optional().map(|v| match v {
156                            Some(value) => AnyValue::Boolean(value),
157                            None => AnyValue::Null,
158                        }))
159                            as Box<dyn ExactSizeIterator<Item = AnyValue<'_>> + '_>
160                    }
161                },
162                _ => Box::new(self.iter()),
163            }
164        }
165    }
166}
167
168pub struct SeriesIter<'a> {
169    arr: &'a dyn Array,
170    dtype: &'a DataType,
171    idx: usize,
172    len: usize,
173}
174
175impl<'a> Iterator for SeriesIter<'a> {
176    type Item = AnyValue<'a>;
177
178    #[inline]
179    fn next(&mut self) -> Option<Self::Item> {
180        let idx = self.idx;
181
182        if idx == self.len {
183            None
184        } else {
185            self.idx += 1;
186            unsafe { Some(arr_to_any_value(self.arr, idx, self.dtype)) }
187        }
188    }
189
190    fn size_hint(&self) -> (usize, Option<usize>) {
191        (self.len, Some(self.len))
192    }
193}
194
195impl ExactSizeIterator for SeriesIter<'_> {}
196
197#[cfg(test)]
198mod test {
199    use crate::prelude::*;
200
201    #[test]
202    fn test_iter() {
203        let a = Series::new("age".into(), [23, 71, 9].as_ref());
204        let _b = a
205            .i32()
206            .unwrap()
207            .into_iter()
208            .map(|opt_v| opt_v.map(|v| v * 2));
209    }
210
211    #[test]
212    fn test_iter_str() {
213        let data = [Some("John"), Some("Doe"), None];
214        let a: Series = data.into_iter().collect();
215        let b = Series::new("".into(), data);
216        assert_eq!(a, b);
217    }
218
219    #[test]
220    fn test_iter_string() {
221        let data = [Some("John".to_string()), Some("Doe".to_string()), None];
222        let a: Series = data.clone().into_iter().collect();
223        let b = Series::new("".into(), data);
224        assert_eq!(a, b);
225    }
226}
polars_core/series/iterator.rs

polars_core/series/
iterator.rs
