CoCalc -- kernels.rs

GitHub Repository: pola-rs/polars
Path: blob/main/crates/polars-time/src/chunkedarray/kernels.rs
⁸⁴⁰⁷ views
1
//! macros that define kernels for extracting
2
//! `week`, `weekday`, `year`, `hour` etc. from primitive arrays.
3
use arrow::array::{BooleanArray, PrimitiveArray};
4
#[cfg(feature = "dtype-time")]
5
use arrow::temporal_conversions::time64ns_to_time_opt;
6
use 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
};
10
use chrono::{Datelike, Timelike};
11

12
use super::super::windows::calendar::*;
13
use super::*;
14

15
trait PolarsIso {
16
    fn week(&self) -> i8;
17
    fn iso_year(&self) -> i32;
18
}
19

20
impl 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

29
impl 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

38
macro_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

55
macro_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

69
macro_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")]
88
to_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")]
97
to_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")]
106
to_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")]
115
to_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")]
123
to_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")]
132
to_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")]
141
to_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")]
150
to_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")]
162
to_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")]
171
to_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")]
180
to_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")]
189
to_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")]
199
to_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")]
209
to_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")]
218
to_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")]
228
to_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")]
238
to_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")]
248
to_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")]
257
to_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")]
265
to_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")]
273
to_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")]
282
to_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")]
292
to_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")]
302
to_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
);
311

312
Product

Resources

Company
