polars_rows_iter/iter_from_column/
iter_from_column_i32.rs

1use super::*;
2use polars::prelude::*;
3
4impl<'a> IterFromColumn<'a> for i32 {
5    type RawInner = i32;
6    fn create_iter(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<i32>> + 'a>> {
7        create_iter(column)
8    }
9
10    #[inline]
11    fn get_value(polars_value: Option<i32>, column_name: &str, _dtype: &DataType) -> PolarsResult<Self>
12    where
13        Self: Sized,
14    {
15        polars_value.ok_or_else(|| <i32 as IterFromColumn<'a>>::unexpected_null_value_error(column_name))
16    }
17}
18
19impl<'a> IterFromColumn<'a> for Option<i32> {
20    type RawInner = i32;
21    fn create_iter(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<i32>> + 'a>> {
22        create_iter(column)
23    }
24
25    #[inline]
26    fn get_value(polars_value: Option<i32>, _column_name: &str, _dtype: &DataType) -> PolarsResult<Self>
27    where
28        Self: Sized,
29    {
30        Ok(polars_value)
31    }
32}
33
34fn create_iter<'a>(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<i32>> + 'a>> {
35    let column_name = column.name().as_str();
36    let iter = match column.dtype() {
37        DataType::Int32 => Box::new(column.i32()?.iter()),
38        DataType::Date => Box::new(column.date()?.phys.iter()),
39        dtype => {
40            return Err(polars_err!(SchemaMismatch: "Cannot get i32 from column '{column_name}' with dtype : {dtype}"))
41        }
42    };
43
44    Ok(iter)
45}
46
47#[cfg(test)]
48mod tests {
49
50    const ROW_COUNT: usize = 64;
51
52    use crate::*;
53    use itertools::{izip, Itertools};
54    use polars::prelude::*;
55    use rand::{rngs::StdRng, SeedableRng};
56    use shared_test_helpers::*;
57
58    create_test_for_chunked_type!(i32_test, i32, i32, DataType::Int32, ROW_COUNT);
59
60    create_test_for_logical_type!(i32_as_date_test, i32, date, DataType::Date, ROW_COUNT);
61
62    #[test]
63    fn i32_as_time_test() {
64        let mut rng = StdRng::seed_from_u64(0);
65        let height = 64;
66        let dtype = DataType::Time;
67
68        let col = create_column("col", dtype.clone(), false, height, &mut rng);
69        let col_opt = create_column("col_opt", dtype, true, height, &mut rng);
70
71        let col_values = col
72            .as_series()
73            .unwrap()
74            .time()
75            .unwrap()
76            .phys
77            .iter()
78            .map(|v| v.unwrap())
79            .collect_vec();
80
81        let col_opt_values = col_opt.as_series().unwrap().time().unwrap().phys.iter().collect_vec();
82
83        let df = DataFrame::new(vec![col, col_opt]).unwrap();
84
85        let col_iter = col_values.iter();
86        let col_opt_iter = col_opt_values.iter();
87
88        let expected_rows = izip!(col_iter, col_opt_iter)
89            .map(|(&col, &col_opt)| TestRow { col, col_opt })
90            .collect_vec();
91
92        #[derive(Debug, FromDataFrameRow, PartialEq)]
93        struct TestRow {
94            col: i64,
95            col_opt: Option<i64>,
96        }
97
98        let rows = df.rows_iter::<TestRow>().unwrap().map(|v| v.unwrap()).collect_vec();
99
100        assert_eq!(rows, expected_rows)
101    }
102}
polars_rows_iter/iter_from_column/iter_from_column_i32.rs

polars_rows_iter/iter_from_column/
iter_from_column_i32.rs
