fix(python): tz-aware filtering (#5603)

pola-rs · Nov 23, 2022 · c0fdc04 · c0fdc04
1 parent c35c3c3
commit c0fdc04
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 1 deletion.
diff --git a/py-polars/polars/internals/lazy_functions.py b/py-polars/polars/internals/lazy_functions.py
@@ -732,7 +732,11 @@ def lit(
     tu: TimeUnit
     if isinstance(value, datetime):
         tu = "us"
-        return lit(_datetime_to_pl_timestamp(value, tu)).cast(Datetime(tu))
+        e = lit(_datetime_to_pl_timestamp(value, tu)).cast(Datetime(tu))
+        if value.tzinfo is not None:
+            return e.dt.tz_localize(str(value.tzinfo))
+        else:
+            return e
 
     elif isinstance(value, timedelta):
         tu = "us"

diff --git a/py-polars/tests/unit/test_datelike.py b/py-polars/tests/unit/test_datelike.py
@@ -1996,3 +1996,44 @@ def test_tz_aware_strftime() -> None:
             "Fri Nov  4 00:00:00 2022",
         ],
     }
+
+
+def test_tz_aware_filter_lit() -> None:
+    start = datetime(1970, 1, 1)
+    stop = datetime(1970, 1, 1, 7)
+    dt = datetime(1970, 1, 1, 6, tzinfo=zoneinfo.ZoneInfo("America/New_York"))
+
+    assert (
+        pl.DataFrame({"date": pl.date_range(start, stop, "1h")})
+        .with_column(pl.col("date").dt.tz_localize("America/New_York").alias("nyc"))
+        .filter(pl.col("nyc") < dt)
+    ).to_dict(False) == {
+        "date": [
+            datetime(1970, 1, 1, 0, 0),
+            datetime(1970, 1, 1, 1, 0),
+            datetime(1970, 1, 1, 2, 0),
+            datetime(1970, 1, 1, 3, 0),
+            datetime(1970, 1, 1, 4, 0),
+            datetime(1970, 1, 1, 5, 0),
+        ],
+        "nyc": [
+            datetime(
+                1970, 1, 1, 0, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+            datetime(
+                1970, 1, 1, 1, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+            datetime(
+                1970, 1, 1, 2, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+            datetime(
+                1970, 1, 1, 3, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+            datetime(
+                1970, 1, 1, 4, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+            datetime(
+                1970, 1, 1, 5, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York")
+            ),
+        ],
+    }