polars_rows_iter/iter_from_column/
iter_from_column_str.rs

1use super::*;
2use iter_from_column_trait::IterFromColumn;
3use polars::prelude::*;
4
5impl<'a> IterFromColumn<'a> for &'a str {
6    type RawInner = &'a str;
7    fn create_iter(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<&'a str>> + 'a>> {
8        create_iter(column)
9    }
10
11    #[inline]
12    fn get_value(polars_value: Option<&'a str>, column_name: &str, _dtype: &DataType) -> PolarsResult<Self>
13    where
14        Self: Sized,
15    {
16        polars_value.ok_or_else(|| <&'a str as IterFromColumn<'a>>::unexpected_null_value_error(column_name))
17    }
18}
19
20impl<'a> IterFromColumn<'a> for Option<&'a str> {
21    type RawInner = &'a str;
22    fn create_iter(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<&'a str>> + 'a>> {
23        create_iter(column)
24    }
25
26    #[inline]
27    fn get_value(polars_value: Option<&'a str>, _column_name: &str, _dtype: &DataType) -> PolarsResult<Self>
28    where
29        Self: Sized,
30    {
31        Ok(polars_value)
32    }
33}
34
35fn create_str_iter<'a>(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<&'a str>> + 'a>> {
36    Ok(Box::new(column.str()?.iter()))
37}
38
39#[cfg(feature = "dtype-categorical")]
40fn create_cat_iter<'a>(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<&'a str>> + 'a>> {
41    Ok(Box::new(column.cat32()?.iter_str()))
42}
43
44#[cfg(feature = "dtype-categorical")]
45fn create_enum_iter<'a>(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<&'a str>> + 'a>> {
46    Ok(Box::new(column.cat8()?.iter_str()))
47}
48
49pub fn create_iter<'a>(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<&'a str>> + 'a>> {
50    let iter = match column.dtype() {
51        DataType::String => create_str_iter(column)?,
52        #[cfg(feature = "dtype-categorical")]
53        DataType::Categorical(_, _) => create_cat_iter(column)?,
54        #[cfg(feature = "dtype-categorical")]
55        DataType::Enum(_, _) => create_enum_iter(column)?,
56        dtype => {
57            let column_name = column.name().as_str();
58            return Err(
59                polars_err!(SchemaMismatch: "Cannot get &str from column '{column_name}' with dtype '{dtype}'.\
60                                             Make sure to enable 'dtype-categorical' feature for 'Categorical' and 'Enum' dtypes."),
61            );
62        }
63    };
64
65    Ok(iter)
66}
67
68#[cfg(test)]
69mod tests {
70    use crate::*;
71    use itertools::{izip, Itertools};
72    use polars::prelude::*;
73    use rand::{rngs::StdRng, SeedableRng};
74    use shared_test_helpers::*;
75
76    const ROW_COUNT: usize = 64;
77
78    #[test]
79    fn str_test() {
80        let mut rng = StdRng::seed_from_u64(0);
81        let height = ROW_COUNT;
82        let dtype = DataType::String;
83
84        let col = create_column("col", dtype.clone(), false, height, &mut rng);
85        let col_opt = create_column("col_opt", dtype, true, height, &mut rng);
86
87        let col_values = col.str().unwrap().iter().map(|v| v.unwrap().to_owned()).collect_vec();
88        let col_opt_values = col_opt
89            .str()
90            .unwrap()
91            .iter()
92            .map(|v| v.map(|s| s.to_owned()))
93            .collect_vec();
94
95        let df = DataFrame::new(vec![col, col_opt]).unwrap();
96
97        let col_iter = col_values.iter();
98        let col_opt_iter = col_opt_values.iter();
99
100        let expected_rows = izip!(col_iter, col_opt_iter)
101            .map(|(col, col_opt)| TestRow {
102                col: col.as_ref(),
103                col_opt: col_opt.as_ref().map(|v| v.as_str()),
104            })
105            .collect_vec();
106
107        #[derive(Debug, FromDataFrameRow, PartialEq)]
108        struct TestRow<'a> {
109            col: &'a str,
110            col_opt: Option<&'a str>,
111        }
112
113        let rows = df.rows_iter::<TestRow>().unwrap().map(|v| v.unwrap()).collect_vec();
114
115        assert_eq!(rows, expected_rows)
116    }
117
118    #[cfg(feature = "dtype-categorical")]
119    #[test]
120    fn cat_test() {
121        let mut rng = StdRng::seed_from_u64(0);
122        let height = ROW_COUNT;
123
124        let cats = Categories::new(PlSmallStr::EMPTY, PlSmallStr::EMPTY, CategoricalPhysical::U32);
125        let dtype = DataType::from_categories(cats);
126
127        let col = create_column("col", dtype.clone(), false, height, &mut rng);
128        let col_opt = create_column("col_opt", dtype, true, height, &mut rng);
129
130        let col_values = col
131            .cat32()
132            .unwrap()
133            .iter_str()
134            .map(|v| v.unwrap().to_owned())
135            .collect_vec();
136        let col_opt_values = col_opt
137            .cat32()
138            .unwrap()
139            .iter_str()
140            .map(|v| v.map(|s| s.to_owned()))
141            .collect_vec();
142
143        let df = DataFrame::new(vec![col, col_opt]).unwrap();
144
145        let col_iter = col_values.iter();
146        let col_opt_iter = col_opt_values.iter();
147
148        let expected_rows = izip!(col_iter, col_opt_iter)
149            .map(|(col, col_opt)| TestRow {
150                col: col.as_ref(),
151                col_opt: col_opt.as_ref().map(|v| v.as_str()),
152            })
153            .collect_vec();
154
155        #[derive(Debug, FromDataFrameRow, PartialEq)]
156        struct TestRow<'a> {
157            col: &'a str,
158            col_opt: Option<&'a str>,
159        }
160
161        let rows = df.rows_iter::<TestRow>().unwrap().map(|v| v.unwrap()).collect_vec();
162
163        assert_eq!(rows, expected_rows)
164    }
165
166    #[cfg(feature = "dtype-categorical")]
167    #[test]
168    fn enum_test() {
169        let mut rng = StdRng::seed_from_u64(0);
170        let height = ROW_COUNT;
171
172        let categories = FrozenCategories::new(["A", "B", "C", "D", "E"]).unwrap();
173        let dtype = DataType::from_frozen_categories(categories);
174
175        let col = create_column("col", dtype.clone(), false, height, &mut rng);
176        let col_opt = create_column("col_opt", dtype, true, height, &mut rng);
177
178        let col_values = col
179            .cat8()
180            .unwrap()
181            .iter_str()
182            .map(|v| v.unwrap().to_owned())
183            .collect_vec();
184        let col_opt_values = col_opt
185            .cat8()
186            .unwrap()
187            .iter_str()
188            .map(|v| v.map(|s| s.to_owned()))
189            .collect_vec();
190
191        let df = DataFrame::new(vec![col, col_opt]).unwrap();
192
193        let col_iter = col_values.iter();
194        let col_opt_iter = col_opt_values.iter();
195
196        let expected_rows = izip!(col_iter, col_opt_iter)
197            .map(|(col, col_opt)| TestRow {
198                col: col.as_ref(),
199                col_opt: col_opt.as_ref().map(|v| v.as_str()),
200            })
201            .collect_vec();
202
203        #[derive(Debug, FromDataFrameRow, PartialEq)]
204        struct TestRow<'a> {
205            col: &'a str,
206            col_opt: Option<&'a str>,
207        }
208
209        let rows = df.rows_iter::<TestRow>().unwrap().map(|v| v.unwrap()).collect_vec();
210
211        assert_eq!(rows, expected_rows)
212    }
213}
polars_rows_iter/iter_from_column/iter_from_column_str.rs

polars_rows_iter/iter_from_column/
iter_from_column_str.rs
