CoCalc -- datetime.rs

GitHub Repository: pola-rs/polars
Path: blob/main/crates/polars-plan/src/plans/aexpr/function_expr/datetime.rs
⁷⁸⁸⁹ views
1
use super::*;
2

3
#[cfg_attr(feature = "ir_serde", derive(serde::Serialize, serde::Deserialize))]
4
#[derive(Clone, PartialEq, Debug, Eq, Hash)]
5
pub enum IRTemporalFunction {
6
    Millennium,
7
    Century,
8
    Year,
9
    IsLeapYear,
10
    IsoYear,
11
    Quarter,
12
    Month,
13
    DaysInMonth,
14
    Week,
15
    WeekDay,
16
    Day,
17
    OrdinalDay,
18
    Time,
19
    Date,
20
    Datetime,
21
    #[cfg(feature = "dtype-duration")]
22
    Duration(TimeUnit),
23
    Hour,
24
    Minute,
25
    Second,
26
    Millisecond,
27
    Microsecond,
28
    Nanosecond,
29
    #[cfg(feature = "dtype-duration")]
30
    TotalDays {
31
        fractional: bool,
32
    },
33
    #[cfg(feature = "dtype-duration")]
34
    TotalHours {
35
        fractional: bool,
36
    },
37
    #[cfg(feature = "dtype-duration")]
38
    TotalMinutes {
39
        fractional: bool,
40
    },
41
    #[cfg(feature = "dtype-duration")]
42
    TotalSeconds {
43
        fractional: bool,
44
    },
45
    #[cfg(feature = "dtype-duration")]
46
    TotalMilliseconds {
47
        fractional: bool,
48
    },
49
    #[cfg(feature = "dtype-duration")]
50
    TotalMicroseconds {
51
        fractional: bool,
52
    },
53
    #[cfg(feature = "dtype-duration")]
54
    TotalNanoseconds {
55
        fractional: bool,
56
    },
57
    ToString(String),
58
    CastTimeUnit(TimeUnit),
59
    WithTimeUnit(TimeUnit),
60
    #[cfg(feature = "timezones")]
61
    ConvertTimeZone(TimeZone),
62
    TimeStamp(TimeUnit),
63
    Truncate,
64
    #[cfg(feature = "offset_by")]
65
    OffsetBy,
66
    #[cfg(feature = "month_start")]
67
    MonthStart,
68
    #[cfg(feature = "month_end")]
69
    MonthEnd,
70
    #[cfg(feature = "timezones")]
71
    BaseUtcOffset,
72
    #[cfg(feature = "timezones")]
73
    DSTOffset,
74
    Round,
75
    Replace,
76
    #[cfg(feature = "timezones")]
77
    ReplaceTimeZone(Option<TimeZone>, NonExistent),
78
    Combine(TimeUnit),
79
    DatetimeFunction {
80
        time_unit: TimeUnit,
81
        time_zone: Option<TimeZone>,
82
    },
83
}
84

85
impl IRTemporalFunction {
86
    pub(super) fn get_field(&self, mapper: FieldsMapper) -> PolarsResult<Field> {
87
        use IRTemporalFunction::*;
88
        match self {
89
            Millennium | Century | Year | IsoYear => mapper.with_dtype(DataType::Int32),
90
            OrdinalDay => mapper.with_dtype(DataType::Int16),
91
            Month | DaysInMonth | Quarter | Week | WeekDay | Day | Hour | Minute | Second => {
92
                mapper.with_dtype(DataType::Int8)
93
            },
94
            Millisecond | Microsecond | Nanosecond => mapper.with_dtype(DataType::Int32),
95
            #[cfg(feature = "dtype-duration")]
96
            TotalDays { fractional }
97
            | TotalHours { fractional }
98
            | TotalMinutes { fractional }
99
            | TotalSeconds { fractional }
100
            | TotalMilliseconds { fractional }
101
            | TotalMicroseconds { fractional }
102
            | TotalNanoseconds { fractional } => {
103
                if *fractional {
104
                    mapper.with_dtype(DataType::Float64)
105
                } else {
106
                    mapper.with_dtype(DataType::Int64)
107
                }
108
            },
109
            ToString(_) => mapper.with_dtype(DataType::String),
110
            WithTimeUnit(tu) | CastTimeUnit(tu) => mapper.try_map_dtype(|dt| match dt {
111
                DataType::Duration(_) => Ok(DataType::Duration(*tu)),
112
                DataType::Datetime(_, tz) => Ok(DataType::Datetime(*tu, tz.clone())),
113
                dtype => polars_bail!(ComputeError: "expected duration or datetime, got {}", dtype),
114
            }),
115
            #[cfg(feature = "timezones")]
116
            ConvertTimeZone(tz) => mapper.try_map_dtype(|dt| match dt {
117
                DataType::Datetime(tu, _) => Ok(DataType::Datetime(*tu, Some(tz.clone()))),
118
                dtype => polars_bail!(ComputeError: "expected Datetime, got {}", dtype),
119
            }),
120
            TimeStamp(_) => mapper.with_dtype(DataType::Int64),
121
            IsLeapYear => mapper.with_dtype(DataType::Boolean),
122
            Time => mapper.with_dtype(DataType::Time),
123
            #[cfg(feature = "dtype-duration")]
124
            Duration(tu) => mapper.with_dtype(DataType::Duration(*tu)),
125
            Date => mapper.with_dtype(DataType::Date),
126
            Datetime => mapper.try_map_dtype(|dt| match dt {
127
                DataType::Datetime(tu, _) => Ok(DataType::Datetime(*tu, None)),
128
                dtype => polars_bail!(ComputeError: "expected Datetime, got {}", dtype),
129
            }),
130
            Truncate => mapper.with_same_dtype(),
131
            #[cfg(feature = "offset_by")]
132
            OffsetBy => mapper.with_same_dtype(),
133
            #[cfg(feature = "month_start")]
134
            MonthStart => mapper.with_same_dtype(),
135
            #[cfg(feature = "month_end")]
136
            MonthEnd => mapper.with_same_dtype(),
137
            #[cfg(feature = "timezones")]
138
            BaseUtcOffset => mapper.with_dtype(DataType::Duration(TimeUnit::Milliseconds)),
139
            #[cfg(feature = "timezones")]
140
            DSTOffset => mapper.with_dtype(DataType::Duration(TimeUnit::Milliseconds)),
141
            Round => mapper.with_same_dtype(),
142
            Replace => mapper.with_same_dtype(),
143
            #[cfg(feature = "timezones")]
144
            ReplaceTimeZone(tz, _non_existent) => mapper.map_datetime_dtype_timezone(tz.as_ref()),
145
            DatetimeFunction {
146
                time_unit,
147
                time_zone,
148
            } => Ok(Field::new(
149
                PlSmallStr::from_static("datetime"),
150
                DataType::Datetime(*time_unit, time_zone.clone()),
151
            )),
152
            Combine(tu) => mapper.try_map_dtype(|dt| match dt {
153
                DataType::Datetime(_, tz) => Ok(DataType::Datetime(*tu, tz.clone())),
154
                DataType::Date => Ok(DataType::Datetime(*tu, None)),
155
                dtype => {
156
                    polars_bail!(ComputeError: "expected Date or Datetime, got {}", dtype)
157
                },
158
            }),
159
        }
160
    }
161

162
    pub fn function_options(&self) -> FunctionOptions {
163
        use IRTemporalFunction as T;
164
        match self {
165
            T::Millennium
166
            | T::Century
167
            | T::Year
168
            | T::IsLeapYear
169
            | T::IsoYear
170
            | T::Quarter
171
            | T::Month
172
            | T::DaysInMonth
173
            | T::Week
174
            | T::WeekDay
175
            | T::Day
176
            | T::OrdinalDay
177
            | T::Time
178
            | T::Date
179
            | T::Datetime
180
            | T::Hour
181
            | T::Minute
182
            | T::Second
183
            | T::Millisecond
184
            | T::Microsecond
185
            | T::Nanosecond
186
            | T::ToString(_)
187
            | T::TimeStamp(_)
188
            | T::CastTimeUnit(_)
189
            | T::WithTimeUnit(_) => FunctionOptions::elementwise(),
190
            #[cfg(feature = "dtype-duration")]
191
            T::TotalDays { .. }
192
            | T::TotalHours { .. }
193
            | T::TotalMinutes { .. }
194
            | T::TotalSeconds { .. }
195
            | T::TotalMilliseconds { .. }
196
            | T::TotalMicroseconds { .. }
197
            | T::TotalNanoseconds { .. } => FunctionOptions::elementwise(),
198
            #[cfg(feature = "timezones")]
199
            T::ConvertTimeZone(_) => FunctionOptions::elementwise(),
200
            #[cfg(feature = "month_start")]
201
            T::MonthStart => FunctionOptions::elementwise(),
202
            #[cfg(feature = "month_end")]
203
            T::MonthEnd => FunctionOptions::elementwise(),
204
            #[cfg(feature = "timezones")]
205
            T::BaseUtcOffset | T::DSTOffset => FunctionOptions::elementwise(),
206
            T::Truncate => FunctionOptions::elementwise(),
207
            #[cfg(feature = "offset_by")]
208
            T::OffsetBy => FunctionOptions::elementwise(),
209
            T::Round => FunctionOptions::elementwise(),
210
            T::Replace => FunctionOptions::elementwise(),
211
            #[cfg(feature = "dtype-duration")]
212
            T::Duration(_) => FunctionOptions::elementwise(),
213
            #[cfg(feature = "timezones")]
214
            T::ReplaceTimeZone(_, _) => FunctionOptions::elementwise(),
215
            T::Combine(_) => FunctionOptions::elementwise(),
216
            T::DatetimeFunction { .. } => {
217
                FunctionOptions::elementwise().with_flags(|f| f | FunctionFlags::ALLOW_RENAME)
218
            },
219
        }
220
    }
221
}
222

223
impl Display for IRTemporalFunction {
224
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
225
        use IRTemporalFunction::*;
226
        let s = match self {
227
            Millennium => "millennium",
228
            Century => "century",
229
            Year => "year",
230
            IsLeapYear => "is_leap_year",
231
            IsoYear => "iso_year",
232
            Quarter => "quarter",
233
            Month => "month",
234
            DaysInMonth => "days_in_month",
235
            Week => "week",
236
            WeekDay => "weekday",
237
            Day => "day",
238
            OrdinalDay => "ordinal_day",
239
            Time => "time",
240
            Date => "date",
241
            Datetime => "datetime",
242
            #[cfg(feature = "dtype-duration")]
243
            Duration(_) => "duration",
244
            Hour => "hour",
245
            Minute => "minute",
246
            Second => "second",
247
            Millisecond => "millisecond",
248
            Microsecond => "microsecond",
249
            Nanosecond => "nanosecond",
250
            #[cfg(feature = "dtype-duration")]
251
            TotalDays { .. } => "total_days",
252
            #[cfg(feature = "dtype-duration")]
253
            TotalHours { .. } => "total_hours",
254
            #[cfg(feature = "dtype-duration")]
255
            TotalMinutes { .. } => "total_minutes",
256
            #[cfg(feature = "dtype-duration")]
257
            TotalSeconds { .. } => "total_seconds",
258
            #[cfg(feature = "dtype-duration")]
259
            TotalMilliseconds { .. } => "total_milliseconds",
260
            #[cfg(feature = "dtype-duration")]
261
            TotalMicroseconds { .. } => "total_microseconds",
262
            #[cfg(feature = "dtype-duration")]
263
            TotalNanoseconds { .. } => "total_nanoseconds",
264
            ToString(_) => "to_string",
265
            #[cfg(feature = "timezones")]
266
            ConvertTimeZone(_) => "convert_time_zone",
267
            CastTimeUnit(_) => "cast_time_unit",
268
            WithTimeUnit(_) => "with_time_unit",
269
            TimeStamp(tu) => return write!(f, "dt.timestamp({tu})"),
270
            Truncate => "truncate",
271
            #[cfg(feature = "offset_by")]
272
            OffsetBy => "offset_by",
273
            #[cfg(feature = "month_start")]
274
            MonthStart => "month_start",
275
            #[cfg(feature = "month_end")]
276
            MonthEnd => "month_end",
277
            #[cfg(feature = "timezones")]
278
            BaseUtcOffset => "base_utc_offset",
279
            #[cfg(feature = "timezones")]
280
            DSTOffset => "dst_offset",
281
            Round => "round",
282
            Replace => "replace",
283
            #[cfg(feature = "timezones")]
284
            ReplaceTimeZone(_, _) => "replace_time_zone",
285
            DatetimeFunction { .. } => return write!(f, "dt.datetime"),
286
            Combine(_) => "combine",
287
        };
288
        write!(f, "dt.{s}")
289
    }
290
}
291

292
Product

Resources

Company
