polars-rows-iter 0.12.0

Library for easy and convenient row iteration of polars dataframes
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

use super::*;
use polars::prelude::*;

impl<'a> IterFromColumn<'a> for i64 {
    type RawInner = i64;
    fn create_iter(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<i64>> + 'a>> {
        create_iter(column)
    }

    #[inline]
    fn get_value(polars_value: Option<i64>, column_name: &str, _dtype: &DataType) -> PolarsResult<Self>
    where
        Self: Sized,
    {
        polars_value.ok_or_else(|| <i64 as IterFromColumn<'a>>::unexpected_null_value_error(column_name))
    }
}

impl<'a> IterFromColumn<'a> for Option<i64> {
    type RawInner = i64;
    fn create_iter(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<i64>> + 'a>> {
        create_iter(column)
    }

    #[inline]
    fn get_value(polars_value: Option<i64>, _column_name: &str, _dtype: &DataType) -> PolarsResult<Self>
    where
        Self: Sized,
    {
        Ok(polars_value)
    }
}

fn create_iter<'a>(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<i64>> + 'a>> {
    let column_name = column.name().as_str();
    let iter = match column.dtype() {
        DataType::Int64 => Box::new(column.i64()?.iter()),
        DataType::Time => Box::new(column.as_materialized_series().time()?.phys.iter()),
        DataType::Datetime(_, _) => Box::new(column.datetime()?.phys.iter()),
        DataType::Duration(_) => Box::new(column.duration()?.phys.iter()),
        dtype => {
            return Err(polars_err!(SchemaMismatch: "Cannot get i64 from column '{column_name}' with dtype : {dtype}"))
        }
    };

    Ok(iter)
}

#[cfg(test)]
mod tests {

    const ROW_COUNT: usize = 64;

    use crate::*;
    use itertools::{izip, Itertools};
    use polars::prelude::*;
    use rand::{rngs::StdRng, SeedableRng};
    use shared_test_helpers::*;

    create_rows_iter_test_for_chunked_type!(i64_test, i64, i64, DataType::Int64, ROW_COUNT);

    create_rows_iter_test_for_logical_type!(
        i64_as_datetime_milliseconds_test,
        i64,
        datetime,
        DataType::Datetime(TimeUnit::Milliseconds, None),
        ROW_COUNT
    );

    create_rows_iter_test_for_logical_type!(
        i64_as_datetime_microseconds_test,
        i64,
        datetime,
        DataType::Datetime(TimeUnit::Microseconds, None),
        ROW_COUNT
    );

    create_rows_iter_test_for_logical_type!(
        i64_as_datetime_nanoseconds_test,
        i64,
        datetime,
        DataType::Datetime(TimeUnit::Nanoseconds, None),
        ROW_COUNT
    );

    #[cfg(feature = "dtype-time")]
    create_rows_iter_test_for_logical_type!(i64_as_time_test, i64, time, DataType::Time, ROW_COUNT);

    create_rows_iter_test_for_logical_type!(
        i64_as_duration_milliseconds_test,
        i64,
        duration,
        DataType::Duration(TimeUnit::Milliseconds),
        ROW_COUNT
    );

    create_rows_iter_test_for_logical_type!(
        i64_as_duration_microseconds_test,
        i64,
        duration,
        DataType::Duration(TimeUnit::Microseconds),
        ROW_COUNT
    );

    create_rows_iter_test_for_logical_type!(
        i64_as_duration_nanoseconds_test,
        i64,
        duration,
        DataType::Duration(TimeUnit::Nanoseconds),
        ROW_COUNT
    );
}