fix(rust, python): set timezone on groupby_dynamic boundaries (#5233)

pola-rs · Oct 17, 2022 · 1ac7b94 · 1ac7b94
1 parent afe8b7c
commit 1ac7b94
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 5 deletions.
diff --git a/polars/polars-time/src/groupby/dynamic.rs b/polars/polars-time/src/groupby/dynamic.rs
@@ -202,6 +202,7 @@ impl Wrap<&DataFrame> {
     ) -> PolarsResult<(Series, Vec<Series>, GroupsProxy)> {
         let w = Window::new(options.every, options.period, options.offset);
         let dt = dt.datetime().unwrap();
+        let tz = dt.time_zone();
 
         let mut lower_bound = None;
         let mut upper_bound = None;
@@ -254,10 +255,10 @@ impl Wrap<&DataFrame> {
                             tu,
                         );
                         let _lower = Int64Chunked::new_vec("lower", lower.clone())
-                            .into_datetime(tu, None)
+                            .into_datetime(tu, tz.clone())
                             .into_series();
                         let _higher = Int64Chunked::new_vec("upper", upper.clone())
-                            .into_datetime(tu, None)
+                            .into_datetime(tu, tz.clone())
                             .into_series();
 
                         update_bounds(lower, upper);
@@ -308,11 +309,11 @@ impl Wrap<&DataFrame> {
             (options.include_boundaries, lower_bound, upper_bound)
         {
             let s = Int64Chunked::new_vec(LB_NAME, lower)
-                .into_datetime(tu, None)
+                .into_datetime(tu, tz.clone())
                 .into_series();
             by.push(s);
             let s = Int64Chunked::new_vec(UP_NAME, higher)
-                .into_datetime(tu, None)
+                .into_datetime(tu, tz.clone())
                 .into_series();
             by.push(s);
         }

diff --git a/py-polars/tests/unit/test_rolling.py b/py-polars/tests/unit/test_rolling.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from datetime import datetime
+import typing
+from datetime import datetime, timedelta
 from typing import TYPE_CHECKING
 
 import pytest
@@ -381,3 +382,30 @@ def test_rolling_var_numerical_stability_5197() -> None:
         0.0,
         0.0,
     ]
+
+
+@typing.no_type_check
+def test_dynamic_groupby_timezone_awareness() -> None:
+    df = pl.DataFrame(
+        (
+            pl.date_range(
+                datetime(2020, 1, 1),
+                datetime(2020, 1, 10),
+                timedelta(days=1),
+                time_unit="ns",
+                name="datetime",
+            ).dt.with_time_zone("UTC"),
+            pl.Series("value", pl.arange(1, 11, eager=True)),
+        )
+    )
+
+    assert (
+        df.groupby_dynamic(
+            "datetime",
+            "3d",
+            offset="-1d",
+            closed="right",
+            include_boundaries=True,
+            truncate=False,
+        ).agg(pl.col("value").last())
+    ).dtypes == [pl.Datetime("ns", "UTC")] * 3 + [pl.Int64]