polars_time/
date_range.rs

1use arrow::legacy::time_zone::Tz;
2use chrono::{Datelike, NaiveDateTime, NaiveTime};
3use polars_core::chunked_array::temporal::time_to_time64ns;
4use polars_core::prelude::*;
5use polars_core::series::IsSorted;
6
7use crate::prelude::*;
8
9pub fn in_nanoseconds_window(ndt: &NaiveDateTime) -> bool {
10    // ~584 year around 1970
11    !(ndt.year() > 2554 || ndt.year() < 1386)
12}
13
14/// Create a [`DatetimeChunked`] from a given `start` and `end` date and a given `interval`.
15pub fn date_range(
16    name: PlSmallStr,
17    start: NaiveDateTime,
18    end: NaiveDateTime,
19    interval: Duration,
20    closed: ClosedWindow,
21    tu: TimeUnit,
22    tz: Option<&Tz>,
23) -> PolarsResult<DatetimeChunked> {
24    let (start, end) = match tu {
25        TimeUnit::Nanoseconds => (
26            start.and_utc().timestamp_nanos_opt().unwrap(),
27            end.and_utc().timestamp_nanos_opt().unwrap(),
28        ),
29        TimeUnit::Microseconds => (
30            start.and_utc().timestamp_micros(),
31            end.and_utc().timestamp_micros(),
32        ),
33        TimeUnit::Milliseconds => (
34            start.and_utc().timestamp_millis(),
35            end.and_utc().timestamp_millis(),
36        ),
37    };
38    datetime_range_impl(name, start, end, interval, closed, tu, tz)
39}
40
41#[doc(hidden)]
42pub fn datetime_range_impl(
43    name: PlSmallStr,
44    start: i64,
45    end: i64,
46    interval: Duration,
47    closed: ClosedWindow,
48    tu: TimeUnit,
49    tz: Option<&Tz>,
50) -> PolarsResult<DatetimeChunked> {
51    let out = Int64Chunked::new_vec(
52        name,
53        datetime_range_i64(start, end, interval, closed, tu, tz)?,
54    );
55    let mut out = match tz {
56        #[cfg(feature = "timezones")]
57        Some(tz) => out.into_datetime(tu, Some(TimeZone::from_chrono(tz))),
58        _ => out.into_datetime(tu, None),
59    };
60
61    out.set_sorted_flag(IsSorted::Ascending);
62    Ok(out)
63}
64
65/// Create a [`TimeChunked`] from a given `start` and `end` date and a given `interval`.
66pub fn time_range(
67    name: PlSmallStr,
68    start: NaiveTime,
69    end: NaiveTime,
70    interval: Duration,
71    closed: ClosedWindow,
72) -> PolarsResult<TimeChunked> {
73    let start = time_to_time64ns(&start);
74    let end = time_to_time64ns(&end);
75    time_range_impl(name, start, end, interval, closed)
76}
77
78#[doc(hidden)]
79pub fn time_range_impl(
80    name: PlSmallStr,
81    start: i64,
82    end: i64,
83    interval: Duration,
84    closed: ClosedWindow,
85) -> PolarsResult<TimeChunked> {
86    let mut out = Int64Chunked::new_vec(
87        name,
88        datetime_range_i64(start, end, interval, closed, TimeUnit::Nanoseconds, None)?,
89    )
90    .into_time();
91
92    out.set_sorted_flag(IsSorted::Ascending);
93    Ok(out)
94}
95
96/// vector of i64 representing temporal values
97pub(crate) fn datetime_range_i64(
98    start: i64,
99    end: i64,
100    interval: Duration,
101    closed: ClosedWindow,
102    time_unit: TimeUnit,
103    time_zone: Option<&Tz>,
104) -> PolarsResult<Vec<i64>> {
105    if start > end {
106        return Ok(Vec::new());
107    }
108    polars_ensure!(
109        !interval.negative && !interval.is_zero(),
110        ComputeError: "`interval` must be positive"
111    );
112
113    let duration = match time_unit {
114        TimeUnit::Nanoseconds => interval.duration_ns(),
115        TimeUnit::Microseconds => interval.duration_us(),
116        TimeUnit::Milliseconds => interval.duration_ms(),
117    };
118    let time_zone_opt: Option<TimeZone> = match time_zone {
119        #[cfg(feature = "timezones")]
120        Some(tz) => Some(TimeZone::from_chrono(tz)),
121        _ => None,
122    };
123
124    if interval.is_constant_duration(time_zone_opt.as_ref()) {
125        // Fast path!
126        let step: usize = duration.try_into().map_err(
127            |_err| polars_err!(ComputeError: "Could not convert {:?} to usize", duration),
128        )?;
129        polars_ensure!(
130            step != 0,
131            InvalidOperation: "interval {} is too small for time unit {} and got rounded down to zero",
132            interval,
133            time_unit,
134        );
135        return match closed {
136            ClosedWindow::Both => Ok((start..=end).step_by(step).collect::<Vec<i64>>()),
137            ClosedWindow::None => Ok((start + duration..end).step_by(step).collect::<Vec<i64>>()),
138            ClosedWindow::Left => Ok((start..end).step_by(step).collect::<Vec<i64>>()),
139            ClosedWindow::Right => Ok((start + duration..=end).step_by(step).collect::<Vec<i64>>()),
140        };
141    }
142
143    let size = ((end - start) / duration + 1) as usize;
144    let offset_fn = match time_unit {
145        TimeUnit::Nanoseconds => Duration::add_ns,
146        TimeUnit::Microseconds => Duration::add_us,
147        TimeUnit::Milliseconds => Duration::add_ms,
148    };
149    let mut ts = Vec::with_capacity(size);
150    let mut i = match closed {
151        ClosedWindow::Both | ClosedWindow::Left => 0,
152        ClosedWindow::Right | ClosedWindow::None => 1,
153    };
154    let mut t = offset_fn(&(interval * i), start, time_zone)?;
155    i += 1;
156    match closed {
157        ClosedWindow::Both | ClosedWindow::Right => {
158            while t <= end {
159                ts.push(t);
160                t = offset_fn(&(interval * i), start, time_zone)?;
161                i += 1;
162            }
163        },
164        ClosedWindow::Left | ClosedWindow::None => {
165            while t < end {
166                ts.push(t);
167                t = offset_fn(&(interval * i), start, time_zone)?;
168                i += 1;
169            }
170        },
171    }
172    debug_assert!(size >= ts.len());
173    Ok(ts)
174}
polars_time/date_range.rs

polars_time/
date_range.rs
