polars_time/chunkedarray/
kernels.rs

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

polars_time/chunkedarray/
kernels.rs
