fix(rust, python): make 'truncate' tz-aware (#5522)

pola-rs · Nov 16, 2022 · 8efb0a0 · 8efb0a0
1 parent 35c53af
commit 8efb0a0
Show file tree

Hide file tree

Showing 4 changed files with 86 additions and 2 deletions.
diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs
@@ -144,6 +144,9 @@ pub(super) fn cast_timezone(s: &Series, tz: &str) -> PolarsResult<Series> {
 pub(super) fn tz_localize(s: &Series, tz: &str) -> PolarsResult<Series> {
     let ca = s.datetime()?.clone();
     match ca.time_zone() {
+        Some(tz) if tz == "UTC" => {
+           Ok(ca.with_time_zone(Some("UTC".into())).into_series())
+        }
         Some(tz) if !tz.is_empty() => {
             Err(PolarsError::ComputeError("Cannot localize a tz-aware datetime. Consider using 'dt.with_time_zone' or 'dt.cast_time_zone'".into()))
         },

diff --git a/polars/polars-time/src/chunkedarray/datetime.rs b/polars/polars-time/src/chunkedarray/datetime.rs
@@ -145,6 +145,13 @@ pub trait DatetimeMethods: AsDatetime {
         .into_datetime(tu, None)
     }
 
+    #[cfg(feature = "timezones")]
+    fn apply_tz_offset(&self, tz: &str) -> PolarsResult<DatetimeChunked> {
+        let ca = self.as_datetime();
+        let keep_tz = ca.time_zone().clone();
+        Ok(self.cast_time_zone(tz)?.with_time_zone(keep_tz))
+    }
+
     #[cfg(feature = "timezones")]
     fn cast_time_zone(&self, tz: &str) -> PolarsResult<DatetimeChunked> {
         use chrono_tz::Tz;

diff --git a/polars/polars-time/src/truncate.rs b/polars/polars-time/src/truncate.rs
@@ -20,8 +20,22 @@ impl PolarsTruncate for DatetimeChunked {
             TimeUnit::Milliseconds => Window::truncate_ms,
         };
 
-        self.apply(|t| func(&w, t))
-            .into_datetime(self.time_unit(), self.time_zone().clone())
+        let out = self
+            .apply(|t| func(&w, t))
+            .into_datetime(self.time_unit(), self.time_zone().clone());
+
+        if self.time_zone().is_some() {
+            #[cfg(feature = "timezones")]
+            {
+                out.apply_tz_offset("UTC").unwrap()
+            }
+            #[cfg(not(feature = "timezones"))]
+            {
+                panic!("activate 'timezones' feature")
+            }
+        } else {
+            out
+        }
     }
 }
 

diff --git a/py-polars/tests/unit/test_datelike.py b/py-polars/tests/unit/test_datelike.py
@@ -1856,3 +1856,63 @@ def test_tz_localize() -> None:
             ),
         ]
     }
+
+
+def test_tz_aware_truncate() -> None:
+    test = pl.DataFrame(
+        {
+            "dt": pl.date_range(
+                low=datetime(2022, 11, 1), high=datetime(2022, 11, 4), interval="12h"
+            ).dt.tz_localize("America/New_York")
+        }
+    )
+    assert test.with_column(pl.col("dt").dt.truncate("1d").alias("trunced")).to_dict(
+        False
+    ) == {
+        "dt": [
+            datetime(
+                2022, 11, 1, 0, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+            datetime(
+                2022, 11, 1, 12, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+            datetime(
+                2022, 11, 2, 0, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+            datetime(
+                2022, 11, 2, 12, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+            datetime(
+                2022, 11, 3, 0, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+            datetime(
+                2022, 11, 3, 12, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+            datetime(
+                2022, 11, 4, 0, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+        ],
+        "trunced": [
+            datetime(
+                2022, 11, 1, 0, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+            datetime(
+                2022, 11, 1, 0, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+            datetime(
+                2022, 11, 2, 0, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+            datetime(
+                2022, 11, 2, 0, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+            datetime(
+                2022, 11, 3, 0, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+            datetime(
+                2022, 11, 3, 0, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+            datetime(
+                2022, 11, 4, 0, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+        ],
+    }