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

polars_time/windows/
duration.rs
