feat(rust, python): make date_range timezone aware (#5234)

polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs

@@ -99,5 +99,5 @@ pub(super) fn timestamp(s: &Series, tu: TimeUnit) -> PolarsResult<Series> {
 #[cfg(feature = "timezones")]
 pub(super) fn cast_timezone(s: &Series, tz: &str) -> PolarsResult<Series> {
     let ca = s.datetime()?;
-    ca.cast_timezone(tz).map(|ca| ca.into_series())
+    ca.cast_time_zone(tz).map(|ca| ca.into_series())
 }

polars/polars-time/Cargo.toml

@@ -25,7 +25,7 @@ dtype-duration = ["polars-core/dtype-duration", "polars-core/temporal"]
 rolling_window = ["polars-core/rolling_window", "dtype-duration"]
 private = []
 fmt = ["polars-core/fmt"]
-timezones = ["chrono-tz"]
+timezones = ["chrono-tz", "dtype-datetime"]
 
 test = [
   "dtype-date",

polars/polars-time/src/chunkedarray/datetime.rs

@@ -146,7 +146,7 @@ pub trait DatetimeMethods: AsDatetime {
     }
 
     #[cfg(feature = "timezones")]
-    fn cast_timezone(&self, tz: &str) -> PolarsResult<DatetimeChunked> {
+    fn cast_time_zone(&self, tz: &str) -> PolarsResult<DatetimeChunked> {
         use chrono_tz::Tz;
         let ca = self.as_datetime();
 

polars/polars-time/src/date_range.rs

@@ -17,9 +17,19 @@ pub fn date_range_impl(
     every: Duration,
     closed: ClosedWindow,
     tu: TimeUnit,
+    _tz: Option<TimeZone>,
 ) -> DatetimeChunked {
     let mut out = Int64Chunked::new_vec(name, date_range_vec(start, stop, every, closed, tu))
         .into_datetime(tu, None);
+
+    #[cfg(feature = "timezones")]
+    if let Some(tz) = _tz {
+        out = out
+            .with_time_zone(Some(tz.clone()))
+            .cast_time_zone("UTC")
+            .unwrap() // ensure we store them as UTC
+            .with_time_zone(Some(tz))
+    }
     out.set_sorted(start > stop);
     out
 }
@@ -32,6 +42,7 @@ pub fn date_range(
     every: Duration,
     closed: ClosedWindow,
     tu: TimeUnit,
+    tz: Option<TimeZone>,
 ) -> DatetimeChunked {
     let (start, stop) = match tu {
         TimeUnit::Nanoseconds => (start.timestamp_nanos(), stop.timestamp_nanos()),
@@ -41,5 +52,5 @@ pub fn date_range(
         ),
         TimeUnit::Milliseconds => (start.timestamp_millis(), stop.timestamp_millis()),
     };
-    date_range_impl(name, start, stop, every, closed, tu)
+    date_range_impl(name, start, stop, every, closed, tu, tz)
 }

polars/polars-time/src/groupby/dynamic.rs

@@ -582,6 +582,7 @@ mod test {
             Duration::parse("30m"),
             ClosedWindow::Both,
             TimeUnit::Milliseconds,
+            None,
         )
         .into_series();
 
@@ -628,6 +629,7 @@ mod test {
             Duration::parse("1h"),
             ClosedWindow::Both,
             TimeUnit::Milliseconds,
+            None,
         )
         .into_series();
         assert_eq!(&upper, &range);
@@ -646,6 +648,7 @@ mod test {
             Duration::parse("1h"),
             ClosedWindow::Both,
             TimeUnit::Milliseconds,
+            None,
         )
         .into_series();
         assert_eq!(&upper, &range);

polars/polars-time/src/upsample.rs

@@ -157,6 +157,7 @@ fn upsample_single_impl(
                         every,
                         ClosedWindow::Both,
                         *tu,
+                        None,
                     )
                     .with_time_zone(tz.clone())
                     .into_series()

polars/tests/it/lazy/expressions/expand.rs

@@ -13,6 +13,7 @@ fn test_expand_datetimes_3042() -> PolarsResult<()> {
         Duration::parse("1w"),
         ClosedWindow::Left,
         TimeUnit::Milliseconds,
+        None,
     )
     .into_series();
 

polars/tests/it/lazy/groupby_dynamic.rs

@@ -18,6 +18,7 @@ fn test_groupby_dynamic_week_bounds() -> PolarsResult<()> {
         Duration::parse("1d"),
         ClosedWindow::Left,
         TimeUnit::Milliseconds,
+        None,
     )
     .into_series();
 

py-polars/polars/internals/functions.py

@@ -181,6 +181,7 @@ def date_range(
     closed: ClosedWindow = "both",
     name: str | None = None,
     time_unit: TimeUnit | None = None,
+    time_zone: str | None = None,
 ) -> pli.Series:
     """
     Create a range of type `Datetime` (or `Date`).
@@ -201,6 +202,8 @@ def date_range(
         Name of the output Series.
     time_unit : {None, 'ns', 'us', 'ms'}
         Set the time unit.
+    time_zone:
+        Optional timezone
 
     Notes
     -----
@@ -256,6 +259,22 @@ def date_range(
     low, low_is_date = _ensure_datetime(low)
     high, high_is_date = _ensure_datetime(high)
 
+    if low.tzinfo is not None or time_zone is not None:
+        if low.tzinfo != high.tzinfo:
+            raise ValueError(
+                "Cannot mix different timezone aware datetimes."
+                f" Got: '{low.tzinfo}' and '{high.tzinfo}'."
+            )
+
+        if time_zone is not None and low.tzinfo is not None:
+            if str(low.tzinfo) != time_zone:
+                raise ValueError(
+                    "Given time_zone is different from that timezone aware datetimes."
+                    f" Given: '{time_zone}', got: '{low.tzinfo}'."
+                )
+        if time_zone is None:
+            time_zone = str(low.tzinfo)
+
     tu: TimeUnit
     if time_unit is not None:
         tu = time_unit
@@ -269,7 +288,9 @@ def date_range(
     if name is None:
         name = ""
 
-    dt_range = pli.wrap_s(_py_date_range(start, stop, interval, closed, name, tu))
+    dt_range = pli.wrap_s(
+        _py_date_range(start, stop, interval, closed, name, tu, time_zone)
+    )
     if (
         low_is_date
         and high_is_date

py-polars/src/lib.rs

@@ -31,7 +31,7 @@ use lazy::ToExprs;
 use mimalloc::MiMalloc;
 use polars::functions::{diag_concat_df, hor_concat_df};
 use polars::prelude::Null;
-use polars_core::datatypes::TimeUnit;
+use polars_core::datatypes::{TimeUnit, TimeZone};
 use polars_core::prelude::{DataFrame, IntoSeries, IDX_DTYPE};
 use polars_core::POOL;
 use pyo3::exceptions::PyValueError;
@@ -453,10 +453,19 @@ fn py_date_range(
     closed: Wrap<ClosedWindow>,
     name: &str,
     tu: Wrap<TimeUnit>,
+    tz: Option<TimeZone>,
 ) -> PySeries {
-    polars::time::date_range_impl(name, start, stop, Duration::parse(every), closed.0, tu.0)
-        .into_series()
-        .into()
+    polars::time::date_range_impl(
+        name,
+        start,
+        stop,
+        Duration::parse(every),
+        closed.0,
+        tu.0,
+        tz,
+    )
+    .into_series()
+    .into()
 }
 
 #[pyfunction]

py-polars/tests/unit/test_datelike.py

@@ -1662,3 +1662,34 @@ def test_auto_infer_time_zone() -> None:
     s = pl.Series([dt])
     assert s.dtype == pl.Datetime("us", "Asia/Shanghai")
     assert s[0] == dt
+
+
+def test_timezone_aware_date_range() -> None:
+    low = datetime(2022, 10, 17, 10, tzinfo=zoneinfo.ZoneInfo("Asia/Shanghai"))
+    high = datetime(2022, 11, 17, 10, tzinfo=zoneinfo.ZoneInfo("Asia/Shanghai"))
+
+    assert pl.date_range(low, high, interval=timedelta(days=5)).to_list() == [
+        datetime(2022, 10, 17, 10, 0, tzinfo=zoneinfo.ZoneInfo(key="Asia/Shanghai")),
+        datetime(2022, 10, 22, 10, 0, tzinfo=zoneinfo.ZoneInfo(key="Asia/Shanghai")),
+        datetime(2022, 10, 27, 10, 0, tzinfo=zoneinfo.ZoneInfo(key="Asia/Shanghai")),
+        datetime(2022, 11, 1, 10, 0, tzinfo=zoneinfo.ZoneInfo(key="Asia/Shanghai")),
+        datetime(2022, 11, 6, 10, 0, tzinfo=zoneinfo.ZoneInfo(key="Asia/Shanghai")),
+        datetime(2022, 11, 11, 10, 0, tzinfo=zoneinfo.ZoneInfo(key="Asia/Shanghai")),
+        datetime(2022, 11, 16, 10, 0, tzinfo=zoneinfo.ZoneInfo(key="Asia/Shanghai")),
+    ]
+
+    with pytest.raises(
+        ValueError,
+        match="Cannot mix different timezone aware datetimes. "
+        "Got: 'Asia/Shanghai' and 'None'",
+    ):
+        pl.date_range(
+            low, high.replace(tzinfo=None), interval=timedelta(days=5), time_zone="UTC"
+        )
+
+    with pytest.raises(
+        ValueError,
+        match="Given time_zone is different from that timezone aware datetimes. "
+        "Given: 'UTC', got: 'Asia/Shanghai'.",
+    ):
+        pl.date_range(low, high, interval=timedelta(days=5), time_zone="UTC")