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;
6use polars_utils::format_pl_smallstr;
7
8use crate::prelude::*;
9
10pub fn in_nanoseconds_window(ndt: &NaiveDateTime) -> bool {
11    // ~584 year around 1970
12    !(ndt.year() > 2554 || ndt.year() < 1386)
13}
14
15/// Create a [`DatetimeChunked`] from a given `start` and `end` date and a given `interval`.
16pub fn date_range(
17    name: PlSmallStr,
18    start: NaiveDateTime,
19    end: NaiveDateTime,
20    interval: Duration,
21    closed: ClosedWindow,
22    tu: TimeUnit,
23    tz: Option<&Tz>,
24) -> PolarsResult<DatetimeChunked> {
25    let (start, end) = match tu {
26        TimeUnit::Nanoseconds => (
27            start.and_utc().timestamp_nanos_opt().unwrap(),
28            end.and_utc().timestamp_nanos_opt().unwrap(),
29        ),
30        TimeUnit::Microseconds => (
31            start.and_utc().timestamp_micros(),
32            end.and_utc().timestamp_micros(),
33        ),
34        TimeUnit::Milliseconds => (
35            start.and_utc().timestamp_millis(),
36            end.and_utc().timestamp_millis(),
37        ),
38    };
39    datetime_range_impl(name, start, end, interval, closed, tu, tz)
40}
41
42#[doc(hidden)]
43pub fn datetime_range_impl(
44    name: PlSmallStr,
45    start: i64,
46    end: i64,
47    interval: Duration,
48    closed: ClosedWindow,
49    tu: TimeUnit,
50    tz: Option<&Tz>,
51) -> PolarsResult<DatetimeChunked> {
52    let out = Int64Chunked::new_vec(
53        name,
54        datetime_range_i64(start, end, interval, closed, tu, tz)?,
55    );
56    let mut out = match tz {
57        #[cfg(feature = "timezones")]
58        Some(tz) => out.into_datetime(tu, Some(format_pl_smallstr!("{}", tz))),
59        _ => out.into_datetime(tu, None),
60    };
61
62    out.set_sorted_flag(IsSorted::Ascending);
63    Ok(out)
64}
65
66/// Create a [`TimeChunked`] from a given `start` and `end` date and a given `interval`.
67pub fn time_range(
68    name: PlSmallStr,
69    start: NaiveTime,
70    end: NaiveTime,
71    interval: Duration,
72    closed: ClosedWindow,
73) -> PolarsResult<TimeChunked> {
74    let start = time_to_time64ns(&start);
75    let end = time_to_time64ns(&end);
76    time_range_impl(name, start, end, interval, closed)
77}
78
79#[doc(hidden)]
80pub fn time_range_impl(
81    name: PlSmallStr,
82    start: i64,
83    end: i64,
84    interval: Duration,
85    closed: ClosedWindow,
86) -> PolarsResult<TimeChunked> {
87    let mut out = Int64Chunked::new_vec(
88        name,
89        datetime_range_i64(start, end, interval, closed, TimeUnit::Nanoseconds, None)?,
90    )
91    .into_time();
92
93    out.set_sorted_flag(IsSorted::Ascending);
94    Ok(out)
95}
96
97/// vector of i64 representing temporal values
98pub(crate) fn datetime_range_i64(
99    start: i64,
100    end: i64,
101    interval: Duration,
102    closed: ClosedWindow,
103    time_unit: TimeUnit,
104    time_zone: Option<&Tz>,
105) -> PolarsResult<Vec<i64>> {
106    if start > end {
107        return Ok(Vec::new());
108    }
109    polars_ensure!(
110        !interval.negative && !interval.is_zero(),
111        ComputeError: "`interval` must be positive"
112    );
113
114    let duration = match time_unit {
115        TimeUnit::Nanoseconds => interval.duration_ns(),
116        TimeUnit::Microseconds => interval.duration_us(),
117        TimeUnit::Milliseconds => interval.duration_ms(),
118    };
119    let time_zone_opt_string: Option<String> = match time_zone {
120        #[cfg(feature = "timezones")]
121        Some(tz) => Some(tz.to_string()),
122        _ => None,
123    };
124    if interval.is_constant_duration(time_zone_opt_string.as_deref()) {
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
