polars_time/windows/
duration.rs

1use std::cmp::Ordering;
2use std::fmt::{Display, Formatter};
3use std::ops::{Mul, Neg};
4
5#[cfg(feature = "timezones")]
6use arrow::legacy::kernels::{Ambiguous, NonExistent};
7use arrow::legacy::time_zone::Tz;
8use arrow::temporal_conversions::{
9    MICROSECONDS, MILLISECONDS, NANOSECONDS, timestamp_ms_to_datetime, timestamp_ns_to_datetime,
10    timestamp_us_to_datetime,
11};
12use chrono::{Datelike, NaiveDate, NaiveDateTime, NaiveTime, Timelike};
13use polars_core::datatypes::DataType;
14use polars_core::prelude::{
15    PolarsResult, TimeZone, datetime_to_timestamp_ms, datetime_to_timestamp_ns,
16    datetime_to_timestamp_us, polars_bail,
17};
18use polars_error::polars_ensure;
19#[cfg(feature = "serde")]
20use serde::{Deserialize, Serialize};
21
22use super::calendar::{
23    NS_DAY, NS_HOUR, NS_MICROSECOND, NS_MILLISECOND, NS_MINUTE, NS_SECOND, NS_WEEK,
24};
25#[cfg(feature = "timezones")]
26use crate::utils::{localize_datetime_opt, try_localize_datetime, unlocalize_datetime};
27use crate::windows::calendar::{DAYS_PER_MONTH, is_leap_year};
28
29#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
30#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
31pub struct Duration {
32    // the number of months for the duration
33    months: i64,
34    // the number of weeks for the duration
35    weeks: i64,
36    // the number of days for the duration
37    days: i64,
38    // the number of nanoseconds for the duration
39    nsecs: i64,
40    // indicates if the duration is negative
41    pub(crate) negative: bool,
42    // indicates if an integer string was passed. e.g. "2i"
43    pub parsed_int: bool,
44}
45
46impl PartialOrd<Self> for Duration {
47    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
48        Some(self.cmp(other))
49    }
50}
51
52impl Ord for Duration {
53    fn cmp(&self, other: &Self) -> Ordering {
54        self.duration_ns().cmp(&other.duration_ns())
55    }
56}
57
58impl Neg for Duration {
59    type Output = Self;
60
61    fn neg(self) -> Self::Output {
62        Self {
63            months: self.months,
64            weeks: self.weeks,
65            days: self.days,
66            nsecs: self.nsecs,
67            negative: !self.negative,
68            parsed_int: self.parsed_int,
69        }
70    }
71}
72
73impl Display for Duration {
74    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
75        if self.is_zero() {
76            return write!(f, "0s");
77        }
78        if self.negative {
79            write!(f, "-")?
80        }
81        if self.months > 0 {
82            write!(f, "{}mo", self.months)?
83        }
84        if self.weeks > 0 {
85            write!(f, "{}w", self.weeks)?
86        }
87        if self.days > 0 {
88            write!(f, "{}d", self.days)?
89        }
90        if self.nsecs > 0 {
91            let secs = self.nsecs / NANOSECONDS;
92            if secs * NANOSECONDS == self.nsecs {
93                write!(f, "{}s", secs)?
94            } else {
95                let us = self.nsecs / 1_000;
96                if us * 1_000 == self.nsecs {
97                    write!(f, "{}us", us)?
98                } else {
99                    write!(f, "{}ns", self.nsecs)?
100                }
101            }
102        }
103        Ok(())
104    }
105}
106
107impl Duration {
108    /// Create a new integer size `Duration`
109    pub fn new(fixed_slots: i64) -> Self {
110        Duration {
111            months: 0,
112            weeks: 0,
113            days: 0,
114            nsecs: fixed_slots.abs(),
115            negative: fixed_slots < 0,
116            parsed_int: true,
117        }
118    }
119
120    /// Parse a string into a `Duration`
121    ///
122    /// Strings are composed of a sequence of number-unit pairs, such as `5d` (5 days). A string may begin with a minus
123    /// sign, in which case it is interpreted as a negative duration. Some examples:
124    ///
125    /// * `"1y"`: 1 year
126    /// * `"-1w2d"`: negative 1 week, 2 days (i.e. -9 days)
127    /// * `"3d12h4m25s"`: 3 days, 12 hours, 4 minutes, and 25 seconds
128    ///
129    /// Aside from a leading minus sign, strings may not contain any characters other than numbers and letters
130    /// (including whitespace).
131    ///
132    /// The available units, in ascending order of magnitude, are as follows:
133    ///
134    /// * `ns`: nanosecond
135    /// * `us`: microsecond
136    /// * `ms`: millisecond
137    /// * `s`:  second
138    /// * `m`:  minute
139    /// * `h`:  hour
140    /// * `d`:  day
141    /// * `w`:  week
142    /// * `mo`: calendar month
143    /// * `q`: calendar quarter
144    /// * `y`:  calendar year
145    /// * `i`:  index value (only for {Int32, Int64} dtypes)
146    ///
147    /// By "calendar day", we mean the corresponding time on the next
148    /// day (which may not be 24 hours, depending on daylight savings).
149    /// Similarly for "calendar week", "calendar month", "calendar quarter",
150    /// and "calendar year".
151    ///
152    /// # Panics
153    /// If the given str is invalid for any reason.
154    pub fn parse(duration: &str) -> Self {
155        Self::try_parse(duration).unwrap()
156    }
157
158    #[doc(hidden)]
159    /// Parse SQL-style "interval" string to Duration. Handles verbose
160    /// units (such as 'year', 'minutes', etc.) and whitespace, as
161    /// well as being case-insensitive.
162    pub fn parse_interval(interval: &str) -> Self {
163        Self::try_parse_interval(interval).unwrap()
164    }
165
166    pub fn try_parse(duration: &str) -> PolarsResult<Self> {
167        Self::_parse(duration, false)
168    }
169
170    pub fn try_parse_interval(interval: &str) -> PolarsResult<Self> {
171        Self::_parse(&interval.to_ascii_lowercase(), true)
172    }
173
174    fn _parse(s: &str, as_interval: bool) -> PolarsResult<Self> {
175        let s = if as_interval { s.trim_start() } else { s };
176
177        let parse_type = if as_interval { "interval" } else { "duration" };
178        let num_minus_signs = s.matches('-').count();
179        if num_minus_signs > 1 {
180            polars_bail!(InvalidOperation: "{} string can only have a single minus sign", parse_type);
181        }
182        if num_minus_signs > 0 {
183            if as_interval {
184                // TODO: intervals need to support per-element minus signs
185                polars_bail!(InvalidOperation: "minus signs are not currently supported in interval strings");
186            } else if !s.starts_with('-') {
187                polars_bail!(InvalidOperation: "only a single minus sign is allowed, at the front of the string");
188            }
189        }
190        let mut months = 0;
191        let mut weeks = 0;
192        let mut days = 0;
193        let mut nsecs = 0;
194
195        let negative = s.starts_with('-');
196        let mut iter = s.char_indices().peekable();
197        let mut start = 0;
198
199        // skip the '-' char
200        if negative {
201            start += 1;
202            iter.next().unwrap();
203        }
204        // permissive whitespace for intervals
205        if as_interval {
206            while let Some((i, ch)) = iter.peek() {
207                if *ch == ' ' {
208                    start = *i + 1;
209                    iter.next();
210                } else {
211                    break;
212                }
213            }
214        }
215        // reserve capacity for the longest valid unit ("microseconds")
216        let mut unit = String::with_capacity(12);
217        let mut parsed_int = false;
218        let mut last_ch_opt: Option<char> = None;
219
220        while let Some((i, mut ch)) = iter.next() {
221            if !ch.is_ascii_digit() {
222                let Ok(n) = s[start..i].parse::<i64>() else {
223                    polars_bail!(InvalidOperation:
224                        "expected leading integer in the {} string, found {}",
225                        parse_type, ch
226                    );
227                };
228
229                loop {
230                    match ch {
231                        c if c.is_ascii_alphabetic() => unit.push(c),
232                        ' ' | ',' if as_interval => {},
233                        _ => break,
234                    }
235                    match iter.next() {
236                        Some((i, ch_)) => {
237                            ch = ch_;
238                            start = i
239                        },
240                        None => break,
241                    }
242                }
243                if unit.is_empty() {
244                    polars_bail!(InvalidOperation:
245                        "expected a unit to follow integer in the {} string '{}'",
246                        parse_type, s
247                    );
248                }
249                match &*unit {
250                    // matches that are allowed for both duration/interval
251                    "ns" => nsecs += n,
252                    "us" => nsecs += n * NS_MICROSECOND,
253                    "ms" => nsecs += n * NS_MILLISECOND,
254                    "s" => nsecs += n * NS_SECOND,
255                    "m" => nsecs += n * NS_MINUTE,
256                    "h" => nsecs += n * NS_HOUR,
257                    "d" => days += n,
258                    "w" => weeks += n,
259                    "mo" => months += n,
260                    "q" => months += n * 3,
261                    "y" => months += n * 12,
262                    "i" => {
263                        nsecs += n;
264                        parsed_int = true;
265                    },
266                    _ if as_interval => match &*unit {
267                        // interval-only (verbose/sql) matches
268                        "nanosecond" | "nanoseconds" => nsecs += n,
269                        "microsecond" | "microseconds" => nsecs += n * NS_MICROSECOND,
270                        "millisecond" | "milliseconds" => nsecs += n * NS_MILLISECOND,
271                        "sec" | "secs" | "second" | "seconds" => nsecs += n * NS_SECOND,
272                        "min" | "mins" | "minute" | "minutes" => nsecs += n * NS_MINUTE,
273                        "hour" | "hours" => nsecs += n * NS_HOUR,
274                        "day" | "days" => days += n,
275                        "week" | "weeks" => weeks += n,
276                        "mon" | "mons" | "month" | "months" => months += n,
277                        "quarter" | "quarters" => months += n * 3,
278                        "year" | "years" => months += n * 12,
279                        _ => {
280                            let valid_units = "'year', 'month', 'quarter', 'week', 'day', 'hour', 'minute', 'second', 'millisecond', 'microsecond', 'nanosecond'";
281                            polars_bail!(InvalidOperation: "unit: '{unit}' not supported; available units include: {} (and their plurals)", valid_units);
282                        },
283                    },
284                    _ => {
285                        polars_bail!(InvalidOperation: "unit: '{unit}' not supported; available units are: 'y', 'mo', 'q', 'w', 'd', 'h', 'm', 's', 'ms', 'us', 'ns'");
286                    },
287                }
288                unit.clear();
289            }
290            last_ch_opt = Some(ch);
291        }
292        if let Some(last_ch) = last_ch_opt {
293            if last_ch.is_ascii_digit() {
294                polars_bail!(InvalidOperation:
295                    "expected a unit to follow integer in the {} string '{}'",
296                    parse_type, s
297                );
298            }
299        };
300
301        Ok(Duration {
302            months: months.abs(),
303            weeks: weeks.abs(),
304            days: days.abs(),
305            nsecs: nsecs.abs(),
306            negative,
307            parsed_int,
308        })
309    }
310
311    fn to_positive(v: i64) -> (bool, i64) {
312        if v < 0 { (true, -v) } else { (false, v) }
313    }
314
315    /// Normalize the duration within the interval.
316    /// It will ensure that the output duration is the smallest positive
317    /// duration that is the equivalent of the current duration.
318    #[allow(dead_code)]
319    pub(crate) fn normalize(&self, interval: &Duration) -> Self {
320        if self.months_only() && interval.months_only() {
321            let mut months = self.months() % interval.months();
322
323            match (self.negative, interval.negative) {
324                (true, true) | (true, false) => months = -months + interval.months(),
325                _ => {},
326            }
327            Duration::from_months(months)
328        } else if self.weeks_only() && interval.weeks_only() {
329            let mut weeks = self.weeks() % interval.weeks();
330
331            match (self.negative, interval.negative) {
332                (true, true) | (true, false) => weeks = -weeks + interval.weeks(),
333                _ => {},
334            }
335            Duration::from_weeks(weeks)
336        } else if self.days_only() && interval.days_only() {
337            let mut days = self.days() % interval.days();
338
339            match (self.negative, interval.negative) {
340                (true, true) | (true, false) => days = -days + interval.days(),
341                _ => {},
342            }
343            Duration::from_days(days)
344        } else {
345            let mut offset = self.duration_ns();
346            if offset == 0 {
347                return *self;
348            }
349            let every = interval.duration_ns();
350
351            if offset < 0 {
352                offset += every * ((offset / -every) + 1)
353            } else {
354                offset -= every * (offset / every)
355            }
356            Duration::from_nsecs(offset)
357        }
358    }
359
360    /// Creates a [`Duration`] that represents a fixed number of nanoseconds.
361    pub(crate) fn from_nsecs(v: i64) -> Self {
362        let (negative, nsecs) = Self::to_positive(v);
363        Self {
364            months: 0,
365            weeks: 0,
366            days: 0,
367            nsecs,
368            negative,
369            parsed_int: false,
370        }
371    }
372
373    /// Creates a [`Duration`] that represents a fixed number of months.
374    pub(crate) fn from_months(v: i64) -> Self {
375        let (negative, months) = Self::to_positive(v);
376        Self {
377            months,
378            weeks: 0,
379            days: 0,
380            nsecs: 0,
381            negative,
382            parsed_int: false,
383        }
384    }
385
386    /// Creates a [`Duration`] that represents a fixed number of weeks.
387    pub(crate) fn from_weeks(v: i64) -> Self {
388        let (negative, weeks) = Self::to_positive(v);
389        Self {
390            months: 0,
391            weeks,
392            days: 0,
393            nsecs: 0,
394            negative,
395            parsed_int: false,
396        }
397    }
398
399    /// Creates a [`Duration`] that represents a fixed number of days.
400    pub(crate) fn from_days(v: i64) -> Self {
401        let (negative, days) = Self::to_positive(v);
402        Self {
403            months: 0,
404            weeks: 0,
405            days,
406            nsecs: 0,
407            negative,
408            parsed_int: false,
409        }
410    }
411
412    /// `true` if zero duration.
413    pub fn is_zero(&self) -> bool {
414        self.months == 0 && self.weeks == 0 && self.days == 0 && self.nsecs == 0
415    }
416
417    pub fn months_only(&self) -> bool {
418        self.months != 0 && self.weeks == 0 && self.days == 0 && self.nsecs == 0
419    }
420
421    pub fn months(&self) -> i64 {
422        self.months
423    }
424
425    pub fn weeks_only(&self) -> bool {
426        self.months == 0 && self.weeks != 0 && self.days == 0 && self.nsecs == 0
427    }
428
429    pub fn weeks(&self) -> i64 {
430        self.weeks
431    }
432
433    pub fn days_only(&self) -> bool {
434        self.months == 0 && self.weeks == 0 && self.days != 0 && self.nsecs == 0
435    }
436
437    pub fn days(&self) -> i64 {
438        self.days
439    }
440
441    /// Returns whether the duration consists of full days.
442    ///
443    /// Note that 24 hours is not considered a full day due to possible
444    /// daylight savings time transitions.
445    pub fn is_full_days(&self) -> bool {
446        self.nsecs == 0
447    }
448
449    pub fn is_constant_duration(&self, time_zone: Option<&TimeZone>) -> bool {
450        if time_zone.is_none() || time_zone == Some(&TimeZone::UTC) {
451            self.months == 0
452        } else {
453            // For non-native, non-UTC time zones, 1 calendar day is not
454            // necessarily 24 hours due to daylight savings time.
455            self.months == 0 && self.weeks == 0 && self.days == 0
456        }
457    }
458
459    /// Returns the nanoseconds from the `Duration` without the weeks or months part.
460    pub fn nanoseconds(&self) -> i64 {
461        self.nsecs
462    }
463
464    /// Returns whether duration is negative.
465    pub fn negative(&self) -> bool {
466        self.negative
467    }
468
469    /// Estimated duration of the window duration. Not a very good one if not a constant duration.
470    #[doc(hidden)]
471    pub const fn duration_ns(&self) -> i64 {
472        self.months * 28 * 24 * 3600 * NANOSECONDS
473            + self.weeks * NS_WEEK
474            + self.days * NS_DAY
475            + self.nsecs
476    }
477
478    #[doc(hidden)]
479    pub const fn duration_us(&self) -> i64 {
480        self.months * 28 * 24 * 3600 * MICROSECONDS
481            + (self.weeks * NS_WEEK / 1000 + self.nsecs / 1000 + self.days * NS_DAY / 1000)
482    }
483
484    #[doc(hidden)]
485    pub const fn duration_ms(&self) -> i64 {
486        self.months * 28 * 24 * 3600 * MILLISECONDS
487            + (self.weeks * NS_WEEK / 1_000_000
488                + self.nsecs / 1_000_000
489                + self.days * NS_DAY / 1_000_000)
490    }
491
492    #[doc(hidden)]
493    fn add_month(ts: NaiveDateTime, n_months: i64, negative: bool) -> NaiveDateTime {
494        let mut months = n_months;
495        if negative {
496            months = -months;
497        }
498
499        // Retrieve the current date and increment the values
500        // based on the number of months
501        let mut year = ts.year();
502        let mut month = ts.month() as i32;
503        let mut day = ts.day();
504        year += (months / 12) as i32;
505        month += (months % 12) as i32;
506
507        // if the month overflowed or underflowed, adjust the year
508        // accordingly. Because we add the modulo for the months
509        // the year will only adjust by one
510        if month > 12 {
511            year += 1;
512            month -= 12;
513        } else if month <= 0 {
514            year -= 1;
515            month += 12;
516        }
517
518        // Normalize the day if we are past the end of the month.
519        let last_day_of_month =
520            DAYS_PER_MONTH[is_leap_year(year) as usize][(month - 1) as usize] as u32;
521
522        if day > last_day_of_month {
523            day = last_day_of_month
524        }
525
526        // Retrieve the original time and construct a data
527        // with the new year, month and day
528        let hour = ts.hour();
529        let minute = ts.minute();
530        let sec = ts.second();
531        let nsec = ts.nanosecond();
532        new_datetime(year, month as u32, day, hour, minute, sec, nsec).expect(
533            "Expected valid datetime, please open an issue at https://github.com/pola-rs/polars/issues"
534        )
535    }
536
537    /// Localize result to given time zone, respecting DST fold of original datetime.
538    /// For example, 2022-11-06 01:30:00 CST truncated by 1 hour becomes 2022-11-06 01:00:00 CST,
539    /// whereas 2022-11-06 01:30:00 CDT truncated by 1 hour becomes 2022-11-06 01:00:00 CDT.
540    ///
541    /// * `original_dt_local` - original datetime, without time zone.
542    ///   E.g. if the original datetime was 2022-11-06 01:30:00 CST, then this would
543    ///   be 2022-11-06 01:30:00.
544    /// * `original_dt_utc` - original datetime converted to UTC. E.g. if the
545    ///   original datetime was 2022-11-06 01:30:00 CST, then this would
546    ///   be 2022-11-06 07:30:00.
547    /// * `result_dt_local` - result, without time zone.
548    #[cfg(feature = "timezones")]
549    fn localize_result(
550        &self,
551        original_dt_local: NaiveDateTime,
552        original_dt_utc: NaiveDateTime,
553        result_dt_local: NaiveDateTime,
554        tz: &Tz,
555    ) -> PolarsResult<NaiveDateTime> {
556        match localize_datetime_opt(result_dt_local, tz, Ambiguous::Raise) {
557            Some(dt) => Ok(dt.expect("we didn't use Ambiguous::Null")),
558            None => {
559                if try_localize_datetime(
560                    original_dt_local,
561                    tz,
562                    Ambiguous::Earliest,
563                    NonExistent::Raise,
564                )?
565                .expect("we didn't use Ambiguous::Null or NonExistent::Null")
566                    == original_dt_utc
567                {
568                    Ok(try_localize_datetime(
569                        result_dt_local,
570                        tz,
571                        Ambiguous::Earliest,
572                        NonExistent::Raise,
573                    )?
574                    .expect("we didn't use Ambiguous::Null or NonExistent::Null"))
575                } else if try_localize_datetime(
576                    original_dt_local,
577                    tz,
578                    Ambiguous::Latest,
579                    NonExistent::Raise,
580                )?
581                .expect("we didn't use Ambiguous::Null or NonExistent::Null")
582                    == original_dt_utc
583                {
584                    Ok(try_localize_datetime(
585                        result_dt_local,
586                        tz,
587                        Ambiguous::Latest,
588                        NonExistent::Raise,
589                    )?
590                    .expect("we didn't use Ambiguous::Null or NonExistent::Null"))
591                } else {
592                    unreachable!()
593                }
594            },
595        }
596    }
597
598    fn truncate_subweekly<G, J>(
599        &self,
600        t: i64,
601        tz: Option<&Tz>,
602        duration: i64,
603        _timestamp_to_datetime: G,
604        _datetime_to_timestamp: J,
605    ) -> PolarsResult<i64>
606    where
607        G: Fn(i64) -> NaiveDateTime,
608        J: Fn(NaiveDateTime) -> i64,
609    {
610        match tz {
611            #[cfg(feature = "timezones")]
612            // for UTC, use fastpath below (same as naive)
613            Some(tz) if tz != &chrono_tz::UTC => {
614                let original_dt_utc = _timestamp_to_datetime(t);
615                let original_dt_local = unlocalize_datetime(original_dt_utc, tz);
616                let t = _datetime_to_timestamp(original_dt_local);
617                let mut remainder = t % duration;
618                if remainder < 0 {
619                    remainder += duration
620                }
621                let result_timestamp = t - remainder;
622                let result_dt_local = _timestamp_to_datetime(result_timestamp);
623                let result_dt_utc =
624                    self.localize_result(original_dt_local, original_dt_utc, result_dt_local, tz)?;
625                Ok(_datetime_to_timestamp(result_dt_utc))
626            },
627            _ => {
628                let mut remainder = t % duration;
629                if remainder < 0 {
630                    remainder += duration
631                }
632                Ok(t - remainder)
633            },
634        }
635    }
636
637    fn truncate_weekly<G, J>(
638        &self,
639        t: i64,
640        tz: Option<&Tz>,
641        _timestamp_to_datetime: G,
642        _datetime_to_timestamp: J,
643        daily_duration: i64,
644    ) -> PolarsResult<i64>
645    where
646        G: Fn(i64) -> NaiveDateTime,
647        J: Fn(NaiveDateTime) -> i64,
648    {
649        let _original_dt_utc: Option<NaiveDateTime>;
650        let _original_dt_local: Option<NaiveDateTime>;
651        let t = match tz {
652            #[cfg(feature = "timezones")]
653            // for UTC, use fastpath below (same as naive)
654            Some(tz) if tz != &chrono_tz::UTC => {
655                _original_dt_utc = Some(_timestamp_to_datetime(t));
656                _original_dt_local = Some(unlocalize_datetime(_original_dt_utc.unwrap(), tz));
657                _datetime_to_timestamp(_original_dt_local.unwrap())
658            },
659            _ => {
660                _original_dt_utc = None;
661                _original_dt_local = None;
662                t
663            },
664        };
665        // If we did
666        //   t - (t % (7 * self.weeks * daily_duration))
667        // then the timestamp would get truncated to the previous Thursday,
668        // because 1970-01-01 (timestamp 0) is a Thursday.
669        // So, we adjust by 4 days to get to Monday.
670        let mut remainder = (t - 4 * daily_duration) % (7 * self.weeks * daily_duration);
671        if remainder < 0 {
672            remainder += 7 * self.weeks * daily_duration
673        }
674        let result_t_local = t - remainder;
675        match tz {
676            #[cfg(feature = "timezones")]
677            // for UTC, use fastpath below (same as naive)
678            Some(tz) if tz != &chrono_tz::UTC => {
679                let result_dt_local = _timestamp_to_datetime(result_t_local);
680                let result_dt_utc = self.localize_result(
681                    _original_dt_local.unwrap(),
682                    _original_dt_utc.unwrap(),
683                    result_dt_local,
684                    tz,
685                )?;
686                Ok(_datetime_to_timestamp(result_dt_utc))
687            },
688            _ => Ok(result_t_local),
689        }
690    }
691    fn truncate_monthly<G, J>(
692        &self,
693        t: i64,
694        tz: Option<&Tz>,
695        timestamp_to_datetime: G,
696        datetime_to_timestamp: J,
697        daily_duration: i64,
698    ) -> PolarsResult<i64>
699    where
700        G: Fn(i64) -> NaiveDateTime,
701        J: Fn(NaiveDateTime) -> i64,
702    {
703        let original_dt_utc;
704        let original_dt_local;
705        let t = match tz {
706            #[cfg(feature = "timezones")]
707            // for UTC, use fastpath below (same as naive)
708            Some(tz) if tz != &chrono_tz::UTC => {
709                original_dt_utc = timestamp_to_datetime(t);
710                original_dt_local = unlocalize_datetime(original_dt_utc, tz);
711                datetime_to_timestamp(original_dt_local)
712            },
713            _ => {
714                original_dt_utc = timestamp_to_datetime(t);
715                original_dt_local = original_dt_utc;
716                datetime_to_timestamp(original_dt_local)
717            },
718        };
719
720        // Remove the time of day from the timestamp
721        // e.g. 2020-01-01 12:34:56 -> 2020-01-01 00:00:00
722        let mut remainder_time = t % daily_duration;
723        if remainder_time < 0 {
724            remainder_time += daily_duration
725        }
726        let t = t - remainder_time;
727
728        // Calculate how many months we need to subtract...
729        let (mut year, mut month) = (
730            original_dt_local.year() as i64,
731            original_dt_local.month() as i64,
732        );
733        let total = ((year - 1970) * 12) + (month - 1);
734        let mut remainder_months = total % self.months;
735        if remainder_months < 0 {
736            remainder_months += self.months
737        }
738
739        // ...and translate that to how many days we need to subtract.
740        let mut _is_leap_year = is_leap_year(year as i32);
741        let mut remainder_days = (original_dt_local.day() - 1) as i64;
742        while remainder_months > 12 {
743            let prev_year_is_leap_year = is_leap_year((year - 1) as i32);
744            let add_extra_day =
745                (_is_leap_year && month > 2) || (prev_year_is_leap_year && month <= 2);
746            remainder_days += 365 + add_extra_day as i64;
747            remainder_months -= 12;
748            year -= 1;
749            _is_leap_year = prev_year_is_leap_year;
750        }
751        while remainder_months > 0 {
752            month -= 1;
753            if month == 0 {
754                year -= 1;
755                _is_leap_year = is_leap_year(year as i32);
756                month = 12;
757            }
758            remainder_days += DAYS_PER_MONTH[_is_leap_year as usize][(month - 1) as usize];
759            remainder_months -= 1;
760        }
761
762        match tz {
763            #[cfg(feature = "timezones")]
764            // for UTC, use fastpath below (same as naive)
765            Some(tz) if tz != &chrono_tz::UTC => {
766                let result_dt_local = timestamp_to_datetime(t - remainder_days * daily_duration);
767                let result_dt_utc =
768                    self.localize_result(original_dt_local, original_dt_utc, result_dt_local, tz)?;
769                Ok(datetime_to_timestamp(result_dt_utc))
770            },
771            _ => Ok(t - remainder_days * daily_duration),
772        }
773    }
774
775    #[inline]
776    pub fn truncate_impl<F, G, J>(
777        &self,
778        t: i64,
779        tz: Option<&Tz>,
780        nsecs_to_unit: F,
781        timestamp_to_datetime: G,
782        datetime_to_timestamp: J,
783    ) -> PolarsResult<i64>
784    where
785        F: Fn(i64) -> i64,
786        G: Fn(i64) -> NaiveDateTime,
787        J: Fn(NaiveDateTime) -> i64,
788    {
789        match (self.months, self.weeks, self.days, self.nsecs) {
790            (0, 0, 0, 0) => polars_bail!(ComputeError: "duration cannot be zero"),
791            // truncate by ns/us/ms
792            (0, 0, 0, _) => {
793                let duration = nsecs_to_unit(self.nsecs);
794                self.truncate_subweekly(
795                    t,
796                    tz,
797                    duration,
798                    timestamp_to_datetime,
799                    datetime_to_timestamp,
800                )
801            },
802            // truncate by days
803            (0, 0, _, 0) => {
804                let duration = self.days * nsecs_to_unit(NS_DAY);
805                self.truncate_subweekly(
806                    t,
807                    tz,
808                    duration,
809                    timestamp_to_datetime,
810                    datetime_to_timestamp,
811                )
812            },
813            // truncate by weeks
814            (0, _, 0, 0) => {
815                let duration = nsecs_to_unit(NS_DAY);
816                self.truncate_weekly(
817                    t,
818                    tz,
819                    timestamp_to_datetime,
820                    datetime_to_timestamp,
821                    duration,
822                )
823            },
824            // truncate by months
825            (_, 0, 0, 0) => {
826                let duration = nsecs_to_unit(NS_DAY);
827                self.truncate_monthly(
828                    t,
829                    tz,
830                    timestamp_to_datetime,
831                    datetime_to_timestamp,
832                    duration,
833                )
834            },
835            _ => {
836                polars_bail!(ComputeError: "duration may not mix month, weeks and nanosecond units")
837            },
838        }
839    }
840
841    // Truncate the given ns timestamp by the window boundary.
842    #[inline]
843    pub fn truncate_ns(&self, t: i64, tz: Option<&Tz>) -> PolarsResult<i64> {
844        self.truncate_impl(
845            t,
846            tz,
847            |nsecs| nsecs,
848            timestamp_ns_to_datetime,
849            datetime_to_timestamp_ns,
850        )
851    }
852
853    // Truncate the given ns timestamp by the window boundary.
854    #[inline]
855    pub fn truncate_us(&self, t: i64, tz: Option<&Tz>) -> PolarsResult<i64> {
856        self.truncate_impl(
857            t,
858            tz,
859            |nsecs| nsecs / 1000,
860            timestamp_us_to_datetime,
861            datetime_to_timestamp_us,
862        )
863    }
864
865    // Truncate the given ms timestamp by the window boundary.
866    #[inline]
867    pub fn truncate_ms(&self, t: i64, tz: Option<&Tz>) -> PolarsResult<i64> {
868        self.truncate_impl(
869            t,
870            tz,
871            |nsecs| nsecs / 1_000_000,
872            timestamp_ms_to_datetime,
873            datetime_to_timestamp_ms,
874        )
875    }
876
877    fn add_impl_month_week_or_day<F, G, J>(
878        &self,
879        mut t: i64,
880        tz: Option<&Tz>,
881        nsecs_to_unit: F,
882        timestamp_to_datetime: G,
883        datetime_to_timestamp: J,
884    ) -> PolarsResult<i64>
885    where
886        F: Fn(i64) -> i64,
887        G: Fn(i64) -> NaiveDateTime,
888        J: Fn(NaiveDateTime) -> i64,
889    {
890        let d = self;
891
892        if d.months > 0 {
893            let ts = match tz {
894                #[cfg(feature = "timezones")]
895                // for UTC, use fastpath below (same as naive)
896                Some(tz) if tz != &chrono_tz::UTC => {
897                    unlocalize_datetime(timestamp_to_datetime(t), tz)
898                },
899                _ => timestamp_to_datetime(t),
900            };
901            let dt = Self::add_month(ts, d.months, d.negative);
902            t = match tz {
903                #[cfg(feature = "timezones")]
904                // for UTC, use fastpath below (same as naive)
905                Some(tz) if tz != &chrono_tz::UTC => datetime_to_timestamp(
906                    try_localize_datetime(dt, tz, Ambiguous::Raise, NonExistent::Raise)?
907                        .expect("we didn't use Ambiguous::Null or NonExistent::Null"),
908                ),
909                _ => datetime_to_timestamp(dt),
910            };
911        }
912
913        if d.weeks > 0 {
914            let t_weeks = nsecs_to_unit(NS_WEEK) * self.weeks;
915            match tz {
916                #[cfg(feature = "timezones")]
917                // for UTC, use fastpath below (same as naive)
918                Some(tz) if tz != &chrono_tz::UTC => {
919                    t = datetime_to_timestamp(unlocalize_datetime(timestamp_to_datetime(t), tz));
920                    t += if d.negative { -t_weeks } else { t_weeks };
921                    t = datetime_to_timestamp(
922                        try_localize_datetime(
923                            timestamp_to_datetime(t),
924                            tz,
925                            Ambiguous::Raise,
926                            NonExistent::Raise,
927                        )?
928                        .expect("we didn't use Ambiguous::Null or NonExistent::Null"),
929                    );
930                },
931                _ => t += if d.negative { -t_weeks } else { t_weeks },
932            };
933        }
934
935        if d.days > 0 {
936            let t_days = nsecs_to_unit(NS_DAY) * self.days;
937            match tz {
938                #[cfg(feature = "timezones")]
939                // for UTC, use fastpath below (same as naive)
940                Some(tz) if tz != &chrono_tz::UTC => {
941                    t = datetime_to_timestamp(unlocalize_datetime(timestamp_to_datetime(t), tz));
942                    t += if d.negative { -t_days } else { t_days };
943                    t = datetime_to_timestamp(
944                        try_localize_datetime(
945                            timestamp_to_datetime(t),
946                            tz,
947                            Ambiguous::Raise,
948                            NonExistent::Raise,
949                        )?
950                        .expect("we didn't use Ambiguous::Null or NonExistent::Null"),
951                    );
952                },
953                _ => t += if d.negative { -t_days } else { t_days },
954            };
955        }
956
957        Ok(t)
958    }
959
960    pub fn add_ns(&self, t: i64, tz: Option<&Tz>) -> PolarsResult<i64> {
961        let d = self;
962        let new_t = self.add_impl_month_week_or_day(
963            t,
964            tz,
965            |nsecs| nsecs,
966            timestamp_ns_to_datetime,
967            datetime_to_timestamp_ns,
968        );
969        let nsecs = if d.negative { -d.nsecs } else { d.nsecs };
970        Ok(new_t? + nsecs)
971    }
972
973    pub fn add_us(&self, t: i64, tz: Option<&Tz>) -> PolarsResult<i64> {
974        let d = self;
975        let new_t = self.add_impl_month_week_or_day(
976            t,
977            tz,
978            |nsecs| nsecs / 1000,
979            timestamp_us_to_datetime,
980            datetime_to_timestamp_us,
981        );
982        let nsecs = if d.negative { -d.nsecs } else { d.nsecs };
983        Ok(new_t? + nsecs / 1_000)
984    }
985
986    pub fn add_ms(&self, t: i64, tz: Option<&Tz>) -> PolarsResult<i64> {
987        let d = self;
988        let new_t = self.add_impl_month_week_or_day(
989            t,
990            tz,
991            |nsecs| nsecs / 1_000_000,
992            timestamp_ms_to_datetime,
993            datetime_to_timestamp_ms,
994        );
995        let nsecs = if d.negative { -d.nsecs } else { d.nsecs };
996        Ok(new_t? + nsecs / 1_000_000)
997    }
998}
999
1000impl Mul<i64> for Duration {
1001    type Output = Self;
1002
1003    fn mul(mut self, mut rhs: i64) -> Self {
1004        if rhs < 0 {
1005            rhs = -rhs;
1006            self.negative = !self.negative
1007        }
1008        self.months *= rhs;
1009        self.weeks *= rhs;
1010        self.days *= rhs;
1011        self.nsecs *= rhs;
1012        self
1013    }
1014}
1015
1016fn new_datetime(
1017    year: i32,
1018    month: u32,
1019    days: u32,
1020    hour: u32,
1021    min: u32,
1022    sec: u32,
1023    nano: u32,
1024) -> Option<NaiveDateTime> {
1025    let date = NaiveDate::from_ymd_opt(year, month, days)?;
1026    let time = NaiveTime::from_hms_nano_opt(hour, min, sec, nano)?;
1027    Some(NaiveDateTime::new(date, time))
1028}
1029
1030pub fn ensure_is_constant_duration(
1031    duration: Duration,
1032    time_zone: Option<&TimeZone>,
1033    variable_name: &str,
1034) -> PolarsResult<()> {
1035    polars_ensure!(duration.is_constant_duration(time_zone),
1036        InvalidOperation: "expected `{}` to be a constant duration \
1037            (i.e. one independent of differing month durations or of daylight savings time), got {}.\n\
1038            \n\
1039            You may want to try:\n\
1040            - using `'730h'` instead of `'1mo'`\n\
1041            - using `'24h'` instead of `'1d'` if your series is time-zone-aware", variable_name, duration);
1042    Ok(())
1043}
1044
1045pub fn ensure_duration_matches_dtype(
1046    duration: Duration,
1047    dtype: &DataType,
1048    variable_name: &str,
1049) -> PolarsResult<()> {
1050    match dtype {
1051        DataType::Int64 | DataType::UInt64 | DataType::Int32 | DataType::UInt32 => {
1052            polars_ensure!(duration.parsed_int || duration.is_zero(),
1053                InvalidOperation: "`{}` duration must be a parsed integer (i.e. use '2i', not '2d') when working with a numeric column", variable_name);
1054        },
1055        DataType::Datetime(_, _) | DataType::Date | DataType::Duration(_) | DataType::Time => {
1056            polars_ensure!(!duration.parsed_int,
1057                InvalidOperation: "`{}` duration may not be a parsed integer (i.e. use '2d', not '2i') when working with a temporal column", variable_name);
1058        },
1059        _ => {
1060            polars_bail!(InvalidOperation: "unsupported data type: {} for temporal/index column, expected UInt64, UInt32, Int64, Int32, Datetime, Date, Duration, or Time", dtype)
1061        },
1062    }
1063    Ok(())
1064}
1065
1066#[cfg(test)]
1067mod test {
1068    use super::*;
1069
1070    #[test]
1071    fn test_parse() {
1072        let out = Duration::parse("1ns");
1073        assert_eq!(out.nsecs, 1);
1074        let out = Duration::parse("1ns1ms");
1075        assert_eq!(out.nsecs, NS_MILLISECOND + 1);
1076        let out = Duration::parse("123ns40ms");
1077        assert_eq!(out.nsecs, 40 * NS_MILLISECOND + 123);
1078        let out = Duration::parse("123ns40ms1w");
1079        assert_eq!(out.nsecs, 40 * NS_MILLISECOND + 123);
1080        assert_eq!(out.duration_ns(), 40 * NS_MILLISECOND + 123 + NS_WEEK);
1081        let out = Duration::parse("-123ns40ms1w");
1082        assert!(out.negative);
1083        let out = Duration::parse("5w");
1084        assert_eq!(out.weeks(), 5);
1085    }
1086
1087    #[test]
1088    fn test_add_ns() {
1089        let t = 1;
1090        let seven_days = Duration::parse("7d");
1091        let one_week = Duration::parse("1w");
1092
1093        // add_ns can only error if a time zone is passed, so it's
1094        // safe to unwrap here
1095        assert_eq!(
1096            seven_days.add_ns(t, None).unwrap(),
1097            one_week.add_ns(t, None).unwrap()
1098        );
1099
1100        let seven_days_negative = Duration::parse("-7d");
1101        let one_week_negative = Duration::parse("-1w");
1102
1103        // add_ns can only error if a time zone is passed, so it's
1104        // safe to unwrap here
1105        assert_eq!(
1106            seven_days_negative.add_ns(t, None).unwrap(),
1107            one_week_negative.add_ns(t, None).unwrap()
1108        );
1109    }
1110
1111    #[test]
1112    fn test_display() {
1113        let duration = Duration::parse("1h");
1114        let expected = "3600s";
1115        assert_eq!(format!("{duration}"), expected);
1116        let duration = Duration::parse("1h5ns");
1117        let expected = "3600000000005ns";
1118        assert_eq!(format!("{duration}"), expected);
1119        let duration = Duration::parse("1h5000ns");
1120        let expected = "3600000005us";
1121        assert_eq!(format!("{duration}"), expected);
1122        let duration = Duration::parse("3mo");
1123        let expected = "3mo";
1124        assert_eq!(format!("{duration}"), expected);
1125        let duration = Duration::parse_interval("4 weeks");
1126        let expected = "4w";
1127        assert_eq!(format!("{duration}"), expected);
1128    }
1129}
polars_time/windows/duration.rs

polars_time/windows/
duration.rs
