fix(rust, python): invert cast_time_zone conversion (#5587)

polars/polars-core/src/chunked_array/temporal/datetime.rs

@@ -44,23 +44,25 @@ impl DatetimeChunked {
 
     #[cfg(feature = "timezones")]
     pub fn apply_tz_offset(&self, tz: &str) -> PolarsResult<DatetimeChunked> {
-        let keep_tz = self.time_zone().clone();
-        Ok(self.cast_time_zone(tz)?.with_time_zone(keep_tz))
+        let keep_tz = self.time_zone().as_deref().unwrap_or("UTC").to_string();
+        self.clone()
+            .with_time_zone(Some(tz.into()))
+            .cast_time_zone(&keep_tz)
     }
 
     #[cfg(feature = "timezones")]
     pub fn cast_time_zone(&self, tz: &str) -> PolarsResult<DatetimeChunked> {
         use chrono_tz::Tz;
 
         if let Some(from) = self.time_zone() {
-            let from: Tz = from.parse().map_err(|_| {
+            let old: Tz = from.parse().map_err(|_| {
                 PolarsError::ComputeError(format!("Could not parse timezone: '{}'", tz).into())
             })?;
-            let to: Tz = tz.parse().map_err(|_| {
+            let new: Tz = tz.parse().map_err(|_| {
                 PolarsError::ComputeError(format!("Could not parse timezone: '{}'", tz).into())
             })?;
             let out =
-                self.apply_kernel(&|arr| cast_timezone(arr, self.time_unit().to_arrow(), from, to));
+                self.apply_kernel(&|arr| cast_timezone(arr, self.time_unit().to_arrow(), new, old));
             Ok(out.into_datetime(self.time_unit(), Some(tz.to_string())))
         } else {
             Err(PolarsError::ComputeError(
@@ -86,11 +88,8 @@ impl DatetimeChunked {
         #[allow(unused_mut)]
         let mut ca = self.clone();
         #[cfg(feature = "timezones")]
-        if let Some(tz) = self.time_zone() {
-            ca = ca
-                .with_time_zone(Some("UTC".into()))
-                .cast_time_zone(tz)
-                .unwrap();
+        if self.time_zone().is_some() {
+            ca = ca.cast_time_zone("UTC").unwrap();
         }
 
         let mut ca: Utf8Chunked = ca.apply_kernel_cast(&|arr| {

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

@@ -151,7 +151,7 @@ pub(super) fn tz_localize(s: &Series, tz: &str) -> PolarsResult<Series> {
             Err(PolarsError::ComputeError("Cannot localize a tz-aware datetime. Consider using 'dt.with_time_zone' or 'dt.cast_time_zone'".into()))
         },
         _ => {
-            Ok(ca.with_time_zone(Some(tz.into())).cast_time_zone("UTC")?.with_time_zone(Some(tz.into())).into_series())
+            Ok(ca.with_time_zone(Some("UTC".to_string())).cast_time_zone(tz)?.into_series())
         }
     }
 }

polars/polars-time/src/date_range.rs

@@ -25,10 +25,9 @@ pub fn date_range_impl(
     #[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))
+            .with_time_zone(Some("UTC".to_string()))
+            .cast_time_zone(&tz)
+            .unwrap()
     }
     out.set_sorted(start > stop);
     out

py-polars/polars/internals/construction.py

@@ -313,15 +313,7 @@ def sequence_to_pyseries(
             if dtype == Datetime and value.tzinfo is not None:
                 py_series = PySeries.new_from_anyvalues(name, values)
                 tz = str(value.tzinfo)
-                return (
-                    pli.wrap_s(py_series)
-                    .dt.with_time_zone(tz)  # first inform polars of tz
-                    .dt.cast_time_zone("UTC")  # ensure we store them as utc
-                    .dt.with_time_zone(
-                        tz
-                    )  # conversion lead to utc tz, which is not correct.
-                    ._s
-                )
+                return pli.wrap_s(py_series).dt.tz_localize(tz)._s
 
             # TODO: use anyvalues here (no need to require pyarrow for this).
             arrow_dtype = dtype_to_arrow_type(dtype)

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-02-29 14:00:00 EST
-            2020-03-31 15:00:00 EDT
-            2020-04-30 15:00:00 EDT
+            2020-03-01 00:00:00 EST
+            2020-04-01 01:00:00 EDT
+            2020-05-01 01:00:00 EDT
         ]
         >>> # Timestamps have changed after cast_time_zone
         >>> date.dt.epoch(tu="s")
         shape: (3,)
         Series: 'NYC' [i64]
         [
-            1583002800
-            1585681200
-            1588273200
+            1583038800
+            1585717200
+            1588309200
         ]
 
         """

py-polars/tests/unit/test_datelike.py

@@ -1719,25 +1719,15 @@ def test_invalid_date_parsing_4898() -> None:
 
 
 def test_cast_timezone() -> None:
-    utc = zoneinfo.ZoneInfo("UTC")
     ny = zoneinfo.ZoneInfo("America/New_York")
     assert pl.DataFrame({"a": [datetime(2022, 9, 25, 14)]}).with_column(
-        pl.col("a")
-        .dt.with_time_zone("America/New_York")
-        .dt.cast_time_zone("UTC")
-        .alias("b")
-    ).to_dict(False) == {
-        "a": [datetime(2022, 9, 25, 14, 0)],
-        "b": [datetime(2022, 9, 25, 18, 0, tzinfo=utc)],
-    }
-    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, 10, 0, tzinfo=ny)],
+        "a": [datetime(2022, 9, 25, 14, 0)],
+        "b": [datetime(2022, 9, 25, 14, 0, tzinfo=ny)],
     }