polars_rows_iter/iter_from_column/
iter_from_column_i64.rs

1use super::*;
2use polars::prelude::*;
3
4impl<'a> IterFromColumn<'a> for i64 {
5    type RawInner = i64;
6    fn create_iter(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<i64>> + 'a>> {
7        create_iter(column)
8    }
9
10    #[inline]
11    fn get_value(polars_value: Option<i64>, column_name: &str, _dtype: &DataType) -> PolarsResult<Self>
12    where
13        Self: Sized,
14    {
15        polars_value.ok_or_else(|| <i64 as IterFromColumn<'a>>::unexpected_null_value_error(column_name))
16    }
17}
18
19impl<'a> IterFromColumn<'a> for Option<i64> {
20    type RawInner = i64;
21    fn create_iter(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<i64>> + 'a>> {
22        create_iter(column)
23    }
24
25    #[inline]
26    fn get_value(polars_value: Option<i64>, _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<i64>> + 'a>> {
35    let column_name = column.name().as_str();
36    let iter = match column.dtype() {
37        DataType::Int64 => Box::new(column.i64()?.iter()),
38        DataType::Time => Box::new(column.as_materialized_series().time()?.phys.iter()),
39        DataType::Datetime(_, _) => Box::new(column.datetime()?.phys.iter()),
40        DataType::Duration(_) => Box::new(column.duration()?.phys.iter()),
41        dtype => {
42            return Err(polars_err!(SchemaMismatch: "Cannot get i64 from column '{column_name}' with dtype : {dtype}"))
43        }
44    };
45
46    Ok(iter)
47}
48
49#[cfg(test)]
50mod tests {
51
52    const ROW_COUNT: usize = 64;
53
54    use crate::*;
55    use itertools::{izip, Itertools};
56    use polars::prelude::*;
57    use rand::{rngs::StdRng, SeedableRng};
58    use shared_test_helpers::*;
59
60    create_test_for_chunked_type!(i64_test, i64, i64, DataType::Int64, ROW_COUNT);
61
62    create_test_for_logical_type!(
63        i64_as_datetime_milliseconds_test,
64        i64,
65        datetime,
66        DataType::Datetime(TimeUnit::Milliseconds, None),
67        ROW_COUNT
68    );
69
70    create_test_for_logical_type!(
71        i64_as_datetime_microseconds_test,
72        i64,
73        datetime,
74        DataType::Datetime(TimeUnit::Microseconds, None),
75        ROW_COUNT
76    );
77
78    create_test_for_logical_type!(
79        i64_as_datetime_nanoseconds_test,
80        i64,
81        datetime,
82        DataType::Datetime(TimeUnit::Nanoseconds, None),
83        ROW_COUNT
84    );
85
86    #[cfg(feature = "dtype-time")]
87    create_test_for_logical_type!(i64_as_time_test, i64, time, DataType::Time, ROW_COUNT);
88
89    create_test_for_logical_type!(
90        i64_as_duration_milliseconds_test,
91        i64,
92        duration,
93        DataType::Duration(TimeUnit::Milliseconds),
94        ROW_COUNT
95    );
96
97    create_test_for_logical_type!(
98        i64_as_duration_microseconds_test,
99        i64,
100        duration,
101        DataType::Duration(TimeUnit::Microseconds),
102        ROW_COUNT
103    );
104
105    create_test_for_logical_type!(
106        i64_as_duration_nanoseconds_test,
107        i64,
108        duration,
109        DataType::Duration(TimeUnit::Nanoseconds),
110        ROW_COUNT
111    );
112}