polars_plan/dsl/functions/
temporal.rs

1use chrono::{Datelike, Timelike};
2
3use super::*;
4
5macro_rules! impl_unit_setter {
6    ($fn_name:ident($field:ident)) => {
7        #[doc = concat!("Set the ", stringify!($field))]
8        pub fn $fn_name(mut self, n: Expr) -> Self {
9            self.$field = n.into();
10            self
11        }
12    };
13}
14
15/// Arguments used by `datetime` in order to produce an [`Expr`] of Datetime
16///
17/// Construct a [`DatetimeArgs`] with `DatetimeArgs::new(y, m, d)`. This will set the other time units to `lit(0)`. You
18/// can then set the other fields with the `with_*` methods, or use `with_hms` to set `hour`, `minute`, and `second` all
19/// at once.
20///
21/// # Examples
22/// ```
23/// use polars_plan::prelude::*;
24/// // construct a DatetimeArgs set to July 20, 1969 at 20:17
25/// let args = DatetimeArgs::new(lit(1969), lit(7), lit(20)).with_hms(lit(20), lit(17), lit(0));
26/// // or
27/// let args = DatetimeArgs::new(lit(1969), lit(7), lit(20)).with_hour(lit(20)).with_minute(lit(17));
28///
29/// // construct a DatetimeArgs using existing columns
30/// let args = DatetimeArgs::new(lit(2023), col("month"), col("day"));
31/// ```
32#[derive(Debug, Clone)]
33pub struct DatetimeArgs {
34    pub year: Expr,
35    pub month: Expr,
36    pub day: Expr,
37    pub hour: Expr,
38    pub minute: Expr,
39    pub second: Expr,
40    pub microsecond: Expr,
41    pub time_unit: TimeUnit,
42    pub time_zone: Option<TimeZone>,
43    pub ambiguous: Expr,
44}
45
46impl Default for DatetimeArgs {
47    fn default() -> Self {
48        Self {
49            year: lit(1970),
50            month: lit(1),
51            day: lit(1),
52            hour: lit(0),
53            minute: lit(0),
54            second: lit(0),
55            microsecond: lit(0),
56            time_unit: TimeUnit::Microseconds,
57            time_zone: None,
58            ambiguous: lit(String::from("raise")),
59        }
60    }
61}
62
63impl DatetimeArgs {
64    /// Construct a new `DatetimeArgs` set to `year`, `month`, `day`
65    ///
66    /// Other fields default to `lit(0)`. Use the `with_*` methods to set them.
67    pub fn new(year: Expr, month: Expr, day: Expr) -> Self {
68        Self {
69            year,
70            month,
71            day,
72            ..Default::default()
73        }
74    }
75
76    /// Set `hour`, `minute`, and `second`
77    ///
78    /// Equivalent to
79    /// ```ignore
80    /// self.with_hour(hour)
81    ///     .with_minute(minute)
82    ///     .with_second(second)
83    /// ```
84    pub fn with_hms(self, hour: Expr, minute: Expr, second: Expr) -> Self {
85        Self {
86            hour,
87            minute,
88            second,
89            ..self
90        }
91    }
92
93    impl_unit_setter!(with_year(year));
94    impl_unit_setter!(with_month(month));
95    impl_unit_setter!(with_day(day));
96    impl_unit_setter!(with_hour(hour));
97    impl_unit_setter!(with_minute(minute));
98    impl_unit_setter!(with_second(second));
99    impl_unit_setter!(with_microsecond(microsecond));
100
101    pub fn with_time_unit(self, time_unit: TimeUnit) -> Self {
102        Self { time_unit, ..self }
103    }
104    #[cfg(feature = "timezones")]
105    pub fn with_time_zone(self, time_zone: Option<TimeZone>) -> Self {
106        Self { time_zone, ..self }
107    }
108    #[cfg(feature = "timezones")]
109    pub fn with_ambiguous(self, ambiguous: Expr) -> Self {
110        Self { ambiguous, ..self }
111    }
112
113    fn all_literal(&self) -> bool {
114        use Expr::*;
115        [
116            &self.year,
117            &self.month,
118            &self.day,
119            &self.hour,
120            &self.minute,
121            &self.second,
122            &self.microsecond,
123        ]
124        .iter()
125        .all(|e| matches!(e, Literal(_)))
126    }
127
128    fn as_literal(&self) -> Option<Expr> {
129        if self.time_zone.is_some() || !self.all_literal() {
130            return None;
131        };
132        let Expr::Literal(lv) = &self.year else {
133            unreachable!()
134        };
135        let year = lv.to_any_value()?.extract()?;
136        let Expr::Literal(lv) = &self.month else {
137            unreachable!()
138        };
139        let month = lv.to_any_value()?.extract()?;
140        let Expr::Literal(lv) = &self.day else {
141            unreachable!()
142        };
143        let day = lv.to_any_value()?.extract()?;
144        let Expr::Literal(lv) = &self.hour else {
145            unreachable!()
146        };
147        let hour = lv.to_any_value()?.extract()?;
148        let Expr::Literal(lv) = &self.minute else {
149            unreachable!()
150        };
151        let minute = lv.to_any_value()?.extract()?;
152        let Expr::Literal(lv) = &self.second else {
153            unreachable!()
154        };
155        let second = lv.to_any_value()?.extract()?;
156        let Expr::Literal(lv) = &self.microsecond else {
157            unreachable!()
158        };
159        let ms: u32 = lv.to_any_value()?.extract()?;
160
161        let dt = chrono::NaiveDateTime::default()
162            .with_year(year)?
163            .with_month(month)?
164            .with_day(day)?
165            .with_hour(hour)?
166            .with_minute(minute)?
167            .with_second(second)?
168            .with_nanosecond(ms * 1000)?;
169
170        let ts = match self.time_unit {
171            TimeUnit::Milliseconds => dt.and_utc().timestamp_millis(),
172            TimeUnit::Microseconds => dt.and_utc().timestamp_micros(),
173            TimeUnit::Nanoseconds => dt.and_utc().timestamp_nanos_opt()?,
174        };
175
176        Some(
177            Expr::Literal(LiteralValue::Scalar(Scalar::new(
178                DataType::Datetime(self.time_unit, None),
179                AnyValue::Datetime(ts, self.time_unit, None),
180            )))
181            .alias(PlSmallStr::from_static("datetime")),
182        )
183    }
184}
185
186/// Construct a column of `Datetime` from the provided [`DatetimeArgs`].
187pub fn datetime(args: DatetimeArgs) -> Expr {
188    if let Some(e) = args.as_literal() {
189        return e;
190    }
191
192    let year = args.year;
193    let month = args.month;
194    let day = args.day;
195    let hour = args.hour;
196    let minute = args.minute;
197    let second = args.second;
198    let microsecond = args.microsecond;
199    let time_unit = args.time_unit;
200    let time_zone = args.time_zone;
201    let ambiguous = args.ambiguous;
202
203    let input = vec![
204        year,
205        month,
206        day,
207        hour,
208        minute,
209        second,
210        microsecond,
211        ambiguous,
212    ];
213
214    Expr::Alias(
215        Arc::new(Expr::Function {
216            input,
217            function: FunctionExpr::TemporalExpr(TemporalFunction::DatetimeFunction {
218                time_unit,
219                time_zone,
220            }),
221            options: FunctionOptions::elementwise()
222                .with_flags(|f| f | FunctionFlags::ALLOW_RENAME)
223                .with_fmt_str("datetime"),
224        }),
225        // TODO: follow left-hand rule in Polars 2.0.
226        PlSmallStr::from_static("datetime"),
227    )
228}
229
230/// Arguments used by `duration` in order to produce an [`Expr`] of [`Duration`]
231///
232/// To construct a [`DurationArgs`], use struct literal syntax with `..Default::default()` to leave unspecified fields at
233/// their default value of `lit(0)`, as demonstrated below.
234///
235/// ```
236/// # use polars_plan::prelude::*;
237/// let args = DurationArgs {
238///     days: lit(5),
239///     hours: col("num_hours"),
240///     minutes: col("num_minutes"),
241///     ..Default::default()  // other fields are lit(0)
242/// };
243/// ```
244/// If you prefer builder syntax, `with_*` methods are also available.
245/// ```
246/// # use polars_plan::prelude::*;
247/// let args = DurationArgs::new().with_weeks(lit(42)).with_hours(lit(84));
248/// ```
249#[derive(Debug, Clone)]
250pub struct DurationArgs {
251    pub weeks: Expr,
252    pub days: Expr,
253    pub hours: Expr,
254    pub minutes: Expr,
255    pub seconds: Expr,
256    pub milliseconds: Expr,
257    pub microseconds: Expr,
258    pub nanoseconds: Expr,
259    pub time_unit: TimeUnit,
260}
261
262impl Default for DurationArgs {
263    fn default() -> Self {
264        Self {
265            weeks: lit(0),
266            days: lit(0),
267            hours: lit(0),
268            minutes: lit(0),
269            seconds: lit(0),
270            milliseconds: lit(0),
271            microseconds: lit(0),
272            nanoseconds: lit(0),
273            time_unit: TimeUnit::Microseconds,
274        }
275    }
276}
277
278impl DurationArgs {
279    /// Create a new [`DurationArgs`] with all fields set to `lit(0)`. Use the `with_*` methods to set the fields.
280    pub fn new() -> Self {
281        Self::default()
282    }
283
284    /// Set `hours`, `minutes`, and `seconds`
285    ///
286    /// Equivalent to:
287    ///
288    /// ```ignore
289    /// self.with_hours(hours)
290    ///     .with_minutes(minutes)
291    ///     .with_seconds(seconds)
292    /// ```
293    pub fn with_hms(self, hours: Expr, minutes: Expr, seconds: Expr) -> Self {
294        Self {
295            hours,
296            minutes,
297            seconds,
298            ..self
299        }
300    }
301
302    /// Set `milliseconds`, `microseconds`, and `nanoseconds`
303    ///
304    /// Equivalent to
305    /// ```ignore
306    /// self.with_milliseconds(milliseconds)
307    ///     .with_microseconds(microseconds)
308    ///     .with_nanoseconds(nanoseconds)
309    /// ```
310    pub fn with_fractional_seconds(
311        self,
312        milliseconds: Expr,
313        microseconds: Expr,
314        nanoseconds: Expr,
315    ) -> Self {
316        Self {
317            milliseconds,
318            microseconds,
319            nanoseconds,
320            ..self
321        }
322    }
323
324    impl_unit_setter!(with_weeks(weeks));
325    impl_unit_setter!(with_days(days));
326    impl_unit_setter!(with_hours(hours));
327    impl_unit_setter!(with_minutes(minutes));
328    impl_unit_setter!(with_seconds(seconds));
329    impl_unit_setter!(with_milliseconds(milliseconds));
330    impl_unit_setter!(with_microseconds(microseconds));
331    impl_unit_setter!(with_nanoseconds(nanoseconds));
332
333    fn all_literal(&self) -> bool {
334        use Expr::*;
335        [
336            &self.weeks,
337            &self.days,
338            &self.hours,
339            &self.seconds,
340            &self.minutes,
341            &self.milliseconds,
342            &self.microseconds,
343            &self.nanoseconds,
344        ]
345        .iter()
346        .all(|e| matches!(e, Literal(_)))
347    }
348
349    fn as_literal(&self) -> Option<Expr> {
350        if !self.all_literal() {
351            return None;
352        };
353        let Expr::Literal(lv) = &self.weeks else {
354            unreachable!()
355        };
356        let weeks = lv.to_any_value()?.extract()?;
357        let Expr::Literal(lv) = &self.days else {
358            unreachable!()
359        };
360        let days = lv.to_any_value()?.extract()?;
361        let Expr::Literal(lv) = &self.hours else {
362            unreachable!()
363        };
364        let hours = lv.to_any_value()?.extract()?;
365        let Expr::Literal(lv) = &self.seconds else {
366            unreachable!()
367        };
368        let seconds = lv.to_any_value()?.extract()?;
369        let Expr::Literal(lv) = &self.minutes else {
370            unreachable!()
371        };
372        let minutes = lv.to_any_value()?.extract()?;
373        let Expr::Literal(lv) = &self.milliseconds else {
374            unreachable!()
375        };
376        let milliseconds = lv.to_any_value()?.extract()?;
377        let Expr::Literal(lv) = &self.microseconds else {
378            unreachable!()
379        };
380        let microseconds = lv.to_any_value()?.extract()?;
381        let Expr::Literal(lv) = &self.nanoseconds else {
382            unreachable!()
383        };
384        let nanoseconds = lv.to_any_value()?.extract()?;
385
386        type D = chrono::Duration;
387        let delta = D::weeks(weeks)
388            + D::days(days)
389            + D::hours(hours)
390            + D::seconds(seconds)
391            + D::minutes(minutes)
392            + D::milliseconds(milliseconds)
393            + D::microseconds(microseconds)
394            + D::nanoseconds(nanoseconds);
395
396        let d = match self.time_unit {
397            TimeUnit::Milliseconds => delta.num_milliseconds(),
398            TimeUnit::Microseconds => delta.num_microseconds()?,
399            TimeUnit::Nanoseconds => delta.num_nanoseconds()?,
400        };
401
402        Some(
403            Expr::Literal(LiteralValue::Scalar(Scalar::new(
404                DataType::Duration(self.time_unit),
405                AnyValue::Duration(d, self.time_unit),
406            )))
407            .alias(PlSmallStr::from_static("duration")),
408        )
409    }
410}
411
412/// Construct a column of [`Duration`] from the provided [`DurationArgs`]
413pub fn duration(args: DurationArgs) -> Expr {
414    if let Some(e) = args.as_literal() {
415        return e;
416    }
417    Expr::Function {
418        input: vec![
419            args.weeks,
420            args.days,
421            args.hours,
422            args.minutes,
423            args.seconds,
424            args.milliseconds,
425            args.microseconds,
426            args.nanoseconds,
427        ],
428        function: FunctionExpr::TemporalExpr(TemporalFunction::Duration(args.time_unit)),
429        options: FunctionOptions::elementwise(),
430    }
431}
polars_plan/dsl/functions/temporal.rs

polars_plan/dsl/functions/
temporal.rs
