polars_time/chunkedarray/
kernels.rs

Help
1//! macros that define kernels for extracting
2//! `week`, `weekday`, `year`, `hour` etc. from primitive arrays.
3use arrow::array::{BooleanArray, PrimitiveArray};
4use arrow::compute::arity::unary;
5#[cfg(feature = "dtype-time")]
6use arrow::temporal_conversions::time64ns_to_time_opt;
7use arrow::temporal_conversions::{
8    date32_to_datetime_opt, timestamp_ms_to_datetime_opt, timestamp_ns_to_datetime_opt,
9    timestamp_us_to_datetime_opt,
10};
11use chrono::{Datelike, Timelike};
12
13use super::super::windows::calendar::*;
14use super::*;
15
16trait PolarsIso {
17    fn week(&self) -> i8;
18    fn iso_year(&self) -> i32;
19}
20
21impl PolarsIso for NaiveDateTime {
22    fn week(&self) -> i8 {
23        self.iso_week().week().try_into().unwrap()
24    }
25    fn iso_year(&self) -> i32 {
26        self.iso_week().year()
27    }
28}
29
30impl PolarsIso for NaiveDate {
31    fn week(&self) -> i8 {
32        self.iso_week().week().try_into().unwrap()
33    }
34    fn iso_year(&self) -> i32 {
35        self.iso_week().year()
36    }
37}
38
39macro_rules! to_temporal_unit {
40    ($name: ident, $chrono_method: ident, $to_datetime_fn: expr,
41    $primitive_in: ty,
42    $primitive_out: ty,
43    $dtype_out:expr) => {
44        pub(crate) fn $name(arr: &PrimitiveArray<$primitive_in>) -> ArrayRef {
45            Box::new(unary(
46                arr,
47                |value| {
48                    $to_datetime_fn(value)
49                        .map(|dt| dt.$chrono_method() as $primitive_out)
50                        .unwrap_or(value as $primitive_out)
51                },
52                $dtype_out,
53            )) as ArrayRef
54        }
55    };
56}
57
58macro_rules! to_boolean_temporal_unit {
59    ($name: ident, $chrono_method: ident, $boolean_method: ident, $to_datetime_fn: expr, $dtype_in: ty) => {
60        pub(crate) fn $name(arr: &PrimitiveArray<$dtype_in>) -> ArrayRef {
61            let values = arr
62                .values()
63                .iter()
64                .map(|value| {
65                    $to_datetime_fn(*value)
66                        .map(|dt| $boolean_method(dt.$chrono_method()))
67                        .unwrap_or(false)
68                })
69                .collect::<Vec<_>>();
70            Box::new(BooleanArray::new(
71                ArrowDataType::Boolean,
72                values.into(),
73                arr.validity().cloned(),
74            ))
75        }
76    };
77}
78
79// Dates
80#[cfg(feature = "dtype-date")]
81to_temporal_unit!(
82    date_to_iso_week,
83    week,
84    date32_to_datetime_opt,
85    i32,
86    i8,
87    ArrowDataType::Int8
88);
89#[cfg(feature = "dtype-date")]
90to_temporal_unit!(
91    date_to_iso_year,
92    iso_year,
93    date32_to_datetime_opt,
94    i32,
95    i32,
96    ArrowDataType::Int32
97);
98#[cfg(feature = "dtype-date")]
99to_temporal_unit!(
100    date_to_year,
101    year,
102    date32_to_datetime_opt,
103    i32,
104    i32,
105    ArrowDataType::Int32
106);
107#[cfg(feature = "dtype-date")]
108to_boolean_temporal_unit!(
109    date_to_is_leap_year,
110    year,
111    is_leap_year,
112    date32_to_datetime_opt,
113    i32
114);
115#[cfg(feature = "dtype-date")]
116to_temporal_unit!(
117    date_to_month,
118    month,
119    date32_to_datetime_opt,
120    i32,
121    i8,
122    ArrowDataType::Int8
123);
124#[cfg(feature = "dtype-date")]
125to_temporal_unit!(
126    date_to_day,
127    day,
128    date32_to_datetime_opt,
129    i32,
130    i8,
131    ArrowDataType::Int8
132);
133#[cfg(feature = "dtype-date")]
134to_temporal_unit!(
135    date_to_ordinal,
136    ordinal,
137    date32_to_datetime_opt,
138    i32,
139    i16,
140    ArrowDataType::Int16
141);
142
143// Times
144#[cfg(feature = "dtype-time")]
145to_temporal_unit!(
146    time_to_hour,
147    hour,
148    time64ns_to_time_opt,
149    i64,
150    i8,
151    ArrowDataType::Int8
152);
153#[cfg(feature = "dtype-time")]
154to_temporal_unit!(
155    time_to_minute,
156    minute,
157    time64ns_to_time_opt,
158    i64,
159    i8,
160    ArrowDataType::Int8
161);
162#[cfg(feature = "dtype-time")]
163to_temporal_unit!(
164    time_to_second,
165    second,
166    time64ns_to_time_opt,
167    i64,
168    i8,
169    ArrowDataType::Int8
170);
171#[cfg(feature = "dtype-time")]
172to_temporal_unit!(
173    time_to_nanosecond,
174    nanosecond,
175    time64ns_to_time_opt,
176    i64,
177    i32,
178    ArrowDataType::Int32
179);
180
181#[cfg(feature = "dtype-datetime")]
182to_temporal_unit!(
183    datetime_to_ordinal_ns,
184    ordinal,
185    timestamp_ns_to_datetime_opt,
186    i64,
187    i16,
188    ArrowDataType::Int16
189);
190
191#[cfg(feature = "dtype-datetime")]
192to_temporal_unit!(
193    datetime_to_ordinal_ms,
194    ordinal,
195    timestamp_ms_to_datetime_opt,
196    i64,
197    i16,
198    ArrowDataType::Int16
199);
200#[cfg(feature = "dtype-datetime")]
201to_temporal_unit!(
202    datetime_to_ordinal_us,
203    ordinal,
204    timestamp_us_to_datetime_opt,
205    i64,
206    i16,
207    ArrowDataType::Int16
208);
209
210#[cfg(feature = "dtype-datetime")]
211to_temporal_unit!(
212    datetime_to_iso_year_ns,
213    iso_year,
214    timestamp_ns_to_datetime_opt,
215    i64,
216    i32,
217    ArrowDataType::Int32
218);
219
220#[cfg(feature = "dtype-datetime")]
221to_temporal_unit!(
222    datetime_to_iso_year_us,
223    iso_year,
224    timestamp_us_to_datetime_opt,
225    i64,
226    i32,
227    ArrowDataType::Int32
228);
229
230#[cfg(feature = "dtype-datetime")]
231to_temporal_unit!(
232    datetime_to_iso_year_ms,
233    iso_year,
234    timestamp_ms_to_datetime_opt,
235    i64,
236    i32,
237    ArrowDataType::Int32
238);
239#[cfg(feature = "dtype-datetime")]
240to_boolean_temporal_unit!(
241    datetime_to_is_leap_year_ns,
242    year,
243    is_leap_year,
244    timestamp_ns_to_datetime_opt,
245    i64
246);
247#[cfg(feature = "dtype-datetime")]
248to_boolean_temporal_unit!(
249    datetime_to_is_leap_year_us,
250    year,
251    is_leap_year,
252    timestamp_us_to_datetime_opt,
253    i64
254);
255#[cfg(feature = "dtype-datetime")]
256to_boolean_temporal_unit!(
257    datetime_to_is_leap_year_ms,
258    year,
259    is_leap_year,
260    timestamp_ms_to_datetime_opt,
261    i64
262);
polars_time/chunkedarray/kernels.rs

polars_time/chunkedarray/
kernels.rs
