polars_rows_iter/iter_from_column/
iter_from_column_string.rs

1use super::iter_from_column_str::create_iter;
2use super::*;
3use iter_from_column_trait::IterFromColumn;
4use polars::prelude::*;
5
6impl<'a> IterFromColumn<'a> for String {
7    type RawInner = &'a str;
8    fn create_iter(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<&'a str>> + 'a>> {
9        create_iter(column)
10    }
11
12    #[inline]
13    fn get_value(polars_value: Option<&'a str>, column_name: &str, _dtype: &DataType) -> PolarsResult<Self>
14    where
15        Self: Sized,
16    {
17        Ok(polars_value
18            .ok_or_else(|| <&'a str as IterFromColumn<'a>>::unexpected_null_value_error(column_name))?
19            .to_string())
20    }
21}
22
23impl<'a> IterFromColumn<'a> for Option<String> {
24    type RawInner = &'a str;
25    fn create_iter(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<&'a str>> + 'a>> {
26        create_iter(column)
27    }
28
29    #[inline]
30    fn get_value(polars_value: Option<&'a str>, _column_name: &str, _dtype: &DataType) -> PolarsResult<Self>
31    where
32        Self: Sized,
33    {
34        Ok(polars_value.map(|s| s.to_string()))
35    }
36}
37
38#[cfg(test)]
39mod tests {
40    use crate::*;
41    use itertools::{izip, Itertools};
42    use polars::prelude::*;
43    use rand::{rngs::StdRng, SeedableRng};
44    use shared_test_helpers::*;
45
46    const ROW_COUNT: usize = 64;
47
48    #[test]
49    fn str_test() {
50        let mut rng = StdRng::seed_from_u64(0);
51        let height = ROW_COUNT;
52        let dtype = DataType::String;
53
54        let col = create_column("col", dtype.clone(), false, height, &mut rng);
55        let col_opt = create_column("col_opt", dtype, true, height, &mut rng);
56
57        let col_values = col.str().unwrap().iter().map(|v| v.unwrap().to_owned()).collect_vec();
58        let col_opt_values = col_opt
59            .str()
60            .unwrap()
61            .iter()
62            .map(|v| v.map(|s| s.to_owned()))
63            .collect_vec();
64
65        let df = DataFrame::new(vec![col, col_opt]).unwrap();
66
67        let col_iter = col_values.into_iter();
68        let col_opt_iter = col_opt_values.into_iter();
69
70        let expected_rows = izip!(col_iter, col_opt_iter)
71            .map(|(col, col_opt)| TestRow { col, col_opt })
72            .collect_vec();
73
74        #[derive(Debug, FromDataFrameRow, PartialEq)]
75        struct TestRow {
76            col: String,
77            col_opt: Option<String>,
78        }
79
80        let rows = df.rows_iter::<TestRow>().unwrap().map(|v| v.unwrap()).collect_vec();
81
82        assert_eq!(rows, expected_rows)
83    }
84
85    #[cfg(feature = "dtype-categorical")]
86    #[test]
87    fn cat_test() {
88        let mut rng = StdRng::seed_from_u64(0);
89        let height = ROW_COUNT;
90        let dtype = DataType::Categorical(None, CategoricalOrdering::Physical);
91
92        let col = create_column("col", dtype.clone(), false, height, &mut rng);
93        let col_opt = create_column("col_opt", dtype, true, height, &mut rng);
94
95        let col_values = col
96            .categorical()
97            .unwrap()
98            .iter_str()
99            .map(|v| v.unwrap().to_owned())
100            .collect_vec();
101        let col_opt_values = col_opt
102            .categorical()
103            .unwrap()
104            .iter_str()
105            .map(|v| v.map(|s| s.to_owned()))
106            .collect_vec();
107
108        let df = DataFrame::new(vec![col, col_opt]).unwrap();
109
110        let col_iter = col_values.into_iter();
111        let col_opt_iter = col_opt_values.into_iter();
112
113        let expected_rows = izip!(col_iter, col_opt_iter)
114            .map(|(col, col_opt)| TestRow { col, col_opt })
115            .collect_vec();
116
117        #[derive(Debug, FromDataFrameRow, PartialEq)]
118        struct TestRow {
119            col: String,
120            col_opt: Option<String>,
121        }
122
123        let rows = df.rows_iter::<TestRow>().unwrap().map(|v| v.unwrap()).collect_vec();
124
125        assert_eq!(rows, expected_rows)
126    }
127
128    #[cfg(feature = "dtype-categorical")]
129    #[test]
130    fn enum_test() {
131        let mut rng = StdRng::seed_from_u64(0);
132        let height = ROW_COUNT;
133
134        let enum_value_series = Series::new("enum".into(), &["A", "B", "C", "D", "E"]);
135        let categories = enum_value_series.str().unwrap().downcast_iter().next().unwrap().clone();
136        let dtype = create_enum_dtype(categories);
137
138        let col = create_column("col", dtype.clone(), false, height, &mut rng);
139        let col_opt = create_column("col_opt", dtype, true, height, &mut rng);
140
141        let col_values = col
142            .categorical()
143            .unwrap()
144            .iter_str()
145            .map(|v| v.unwrap().to_owned())
146            .collect_vec();
147        let col_opt_values = col_opt
148            .categorical()
149            .unwrap()
150            .iter_str()
151            .map(|v| v.map(|s| s.to_owned()))
152            .collect_vec();
153
154        let df = DataFrame::new(vec![col, col_opt]).unwrap();
155
156        let col_iter = col_values.into_iter();
157        let col_opt_iter = col_opt_values.into_iter();
158
159        let expected_rows = izip!(col_iter, col_opt_iter)
160            .map(|(col, col_opt)| TestRow { col, col_opt })
161            .collect_vec();
162
163        #[derive(Debug, FromDataFrameRow, PartialEq)]
164        struct TestRow {
165            col: String,
166            col_opt: Option<String>,
167        }
168
169        let rows = df.rows_iter::<TestRow>().unwrap().map(|v| v.unwrap()).collect_vec();
170
171        assert_eq!(rows, expected_rows)
172    }
173}
polars_rows_iter/iter_from_column/iter_from_column_string.rs

polars_rows_iter/iter_from_column/
iter_from_column_string.rs
