fix[rust]: fix timezone cast (#5016)

pola-rs · Sep 28, 2022 · 2cf3d7b · 2cf3d7b
1 parent ee883e0
commit 2cf3d7b
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 9 deletions.
diff --git a/polars/polars-time/src/chunkedarray/kernels.rs b/polars/polars-time/src/chunkedarray/kernels.rs
@@ -234,35 +234,66 @@ pub(crate) fn cast_timezone(
     to: chrono_tz::Tz,
 ) -> ArrayRef {
     use chrono::TimeZone;
+    use chrono_tz::OffsetComponents;
 
     match tu {
         TimeUnit::Milliseconds => Box::new(unary(
             arr,
             |value| {
                 let ndt = timestamp_ms_to_datetime(value);
+                // find the current offset from utc
                 let tz_aware = from.from_local_datetime(&ndt).unwrap();
+                let offset = tz_aware.offset();
+                let total_offset_from = offset.base_utc_offset() + offset.dst_offset();
+
+                // find the new offset from utc
                 let new_tz_aware = tz_aware.with_timezone(&to);
-                new_tz_aware.timestamp_millis()
+                let offset = new_tz_aware.offset();
+                let total_offset_to = offset.base_utc_offset() + offset.dst_offset();
+                let offset = total_offset_to - total_offset_from;
+
+                // correct for that offset
+                (ndt + offset).timestamp_millis()
             },
             ArrowDataType::Int64,
         )),
         TimeUnit::Microseconds => Box::new(unary(
             arr,
             |value| {
                 let ndt = timestamp_us_to_datetime(value);
+                // find the current offset from utc
                 let tz_aware = from.from_local_datetime(&ndt).unwrap();
+                let offset = tz_aware.offset();
+                let total_offset_from = offset.base_utc_offset() + offset.dst_offset();
+
+                // find the new offset from utc
                 let new_tz_aware = tz_aware.with_timezone(&to);
-                new_tz_aware.timestamp_micros()
+                let offset = new_tz_aware.offset();
+                let total_offset_to = offset.base_utc_offset() + offset.dst_offset();
+                let offset = total_offset_to - total_offset_from;
+
+                // correct for that offset
+                (ndt + offset).timestamp_micros()
             },
             ArrowDataType::Int64,
         )),
         TimeUnit::Nanoseconds => Box::new(unary(
             arr,
             |value| {
                 let ndt = timestamp_ns_to_datetime(value);
+                // find the current offset from utc
                 let tz_aware = from.from_local_datetime(&ndt).unwrap();
+                let offset = tz_aware.offset();
+                let total_offset_from = offset.base_utc_offset() + offset.dst_offset();
+
+                // find the new offset from utc
                 let new_tz_aware = tz_aware.with_timezone(&to);
-                new_tz_aware.timestamp_nanos()
+                let offset = new_tz_aware.offset();
+                let total_offset_to = offset.base_utc_offset() + offset.dst_offset();
+                let offset = total_offset_to - total_offset_from;
+
+                // correct for that offset
+                (ndt + offset).timestamp_nanos()
             },
             ArrowDataType::Int64,
         )),

diff --git a/py-polars/polars/internals/series/datetime.py b/py-polars/polars/internals/series/datetime.py
@@ -998,18 +998,18 @@ def cast_time_zone(self, tz: str) -> pli.Series:
         shape: (3,)
         Series: 'NYC' [datetime[μs, America/New_York]]
         [
-                2020-03-01 00:00:00 EST
-                2020-03-31 23:00:00 EDT
-                2020-04-30 23:00:00 EDT
+            2020-02-29 19:00:00 EST
+            2020-03-31 19:00:00 EDT
+            2020-04-30 19:00:00 EDT
         ]
         >>> # Timestamps have changed after cast_time_zone
         >>> date.dt.epoch(tu="s")
         shape: (3,)
         Series: 'NYC' [i64]
         [
-                1583020800
-                1585695600
-                1588287600
+            1583002800
+            1585681200
+            1588273200
         ]
 
         """

diff --git a/py-polars/tests/unit/test_datelike.py b/py-polars/tests/unit/test_datelike.py
@@ -1534,3 +1534,12 @@ def test_cast_timezone() -> None:
         "a": [datetime(2022, 9, 25, 14, 0)],
         "b": [datetime(2022, 9, 25, 18, 0)],
     }
+    assert pl.DataFrame({"a": [datetime(2022, 9, 25, 18)]}).with_column(
+        pl.col("a")
+        .dt.with_time_zone("UTC")
+        .dt.cast_time_zone("America/New_York")
+        .alias("b")
+    ).to_dict(False) == {
+        "a": [datetime(2022, 9, 25, 18, 0)],
+        "b": [datetime(2022, 9, 25, 14, 0)],
+    }