From 79221a8426d556050efa79a127070adf79cef4e5 Mon Sep 17 00:00:00 2001
From: Marco Edward Gorelli <33491632+MarcoGorelli@users.noreply.github.com>
Date: Sun, 19 Feb 2023 16:05:02 +0000
Subject: [PATCH] fix(rust, python): reflect time zone conversion in lazy
 dataframe schema (#7022)

Co-authored-by: MarcoGorelli <>
---
 polars/polars-lazy/polars-plan/src/dsl/dt.rs    |  6 +++++-
 py-polars/tests/unit/datatypes/test_temporal.py | 17 +++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/polars/polars-lazy/polars-plan/src/dsl/dt.rs b/polars/polars-lazy/polars-plan/src/dsl/dt.rs
index 6bd1614f1dd0..f00ff73fddb4 100644
--- a/polars/polars-lazy/polars-plan/src/dsl/dt.rs
+++ b/polars/polars-lazy/polars-plan/src/dsl/dt.rs
@@ -68,6 +68,7 @@ impl DateLikeNameSpace {
     /// Change the underlying [`TimeZone`] of the [`Series`]. This does not modify the data.
     #[cfg(feature = "timezones")]
     pub fn convert_time_zone(self, time_zone: TimeZone) -> Expr {
+        let time_zone_clone = time_zone.clone();
         self.0.map(
             move |s| match s.dtype() {
                 DataType::Datetime(_, Some(_)) => {
@@ -79,7 +80,10 @@ impl DateLikeNameSpace {
                     "Cannot call convert_time_zone on tz-naive. Set a time zone first with replace_time_zone".into()
                 )),
             },
-            GetOutput::same_type(),
+            GetOutput::map_dtype(move |dtype| match dtype {
+                DataType::Datetime(tu, _) => DataType::Datetime(*tu, Some(time_zone_clone.clone())),
+                _ => panic!("expected datetime"),
+            }),
         )
     }
 
diff --git a/py-polars/tests/unit/datatypes/test_temporal.py b/py-polars/tests/unit/datatypes/test_temporal.py
index 6ed5178bd5a2..f172a36e7db2 100644
--- a/py-polars/tests/unit/datatypes/test_temporal.py
+++ b/py-polars/tests/unit/datatypes/test_temporal.py
@@ -1768,6 +1768,23 @@ def test_convert_time_zone_invalid() -> None:
         ts.dt.replace_time_zone("UTC").dt.convert_time_zone("foo")
 
 
+def test_convert_time_zone_lazy_schema() -> None:
+    ts_us = pl.Series(["2020-01-01"]).str.strptime(pl.Datetime("us", "UTC"))
+    ts_ms = pl.Series(["2020-01-01"]).str.strptime(pl.Datetime("ms", "UTC"))
+    ldf = pl.DataFrame({"ts_us": ts_us, "ts_ms": ts_ms}).lazy()
+    result = ldf.with_columns(
+        pl.col("ts_us").dt.convert_time_zone("America/New_York").alias("ts_us_ny"),
+        pl.col("ts_ms").dt.convert_time_zone("America/New_York").alias("ts_us_kt"),
+    ).schema
+    expected = {
+        "ts_us": pl.Datetime("us", "UTC"),
+        "ts_ms": pl.Datetime("ms", "UTC"),
+        "ts_us_ny": pl.Datetime("us", "America/New_York"),
+        "ts_us_kt": pl.Datetime("ms", "America/New_York"),
+    }
+    assert result == expected
+
+
 def test_convert_time_zone_on_tz_naive() -> None:
     ts = pl.Series(["2020-01-01"]).str.strptime(pl.Datetime)
     with pytest.raises(