fix(rust, python): reflect time zone conversion in lazy dataframe sch…

…ema (pola-rs#7022)

Co-authored-by: MarcoGorelli <>

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"),
+            }),
         )
     }
 

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(