Merge pull request #1 from Point72/fix/dt-date-pushdown

Recognize .dt.date() in temporal predicate pushdown

Author: ptomecek

polars_io_tools/io_sources/base.py

@@ -15,6 +15,7 @@
     FunctionType,
     GenericFunctionType,
     OperatorType,
+    TemporalFunctionType,
     TimeUnit,
     get_function_enum,
 )
@@ -1101,6 +1102,10 @@ def extract_column_name(node: BaseExprNode) -> Optional[str]:
             return col_names[0]
         elif isinstance(cur, (AliasNode, CastNode)):
             cur = cur.input
+        elif isinstance(cur, FunctionNode) and cur.function_type == TemporalFunctionType.DATE and len(cur.inputs) == 1:
+            # ``col.dt.date()`` floors a Datetime to day granularity: order- and domain-preserving, hence equivalent
+            # to ``col.cast(pl.Date)`` (unwrapped above). Treat it identically so it resolves wherever a cast would.
+            cur = cur.inputs[0]
         else:
             return None
 

polars_io_tools/tests/io_sources/test_base.py

@@ -79,6 +79,13 @@ def test_extract_column_name():
     expr = pl.col("b").cast(pl.Int64).alias("first").alias("second")
     assert "b" == extract_column_name(get_parsed_expr(expr))
 
+    # .dt.date() floors a Datetime column to day granularity (order-preserving, like cast(pl.Date)) and should
+    # resolve to the underlying column so temporal predicates push down.
+    expr = pl.col("ts").dt.date()
+    assert "ts" == extract_column_name(get_parsed_expr(expr))
+    expr = pl.col("ts").dt.date().alias("d")
+    assert "ts" == extract_column_name(get_parsed_expr(expr))
+
 
 def test_parse_simple_column():
     """Test parsing a simple column reference."""

polars_io_tools/tests/io_sources/test_multi_source.py

@@ -233,6 +233,29 @@ def test_simple_date_filter(self):
         # With maintain_order="left", row order should be preserved
         assert_frame_equal(result, expected)
 
+    def test_dt_date_filter_prunes_datetime_source(self):
+        """A ``.dt.date()`` window on a Datetime source prunes upstream (no full scan) and stays correct.
+
+        ``.dt.date()`` floors to day granularity (order-preserving, like ``cast(pl.Date)``), so the temporal range
+        must still be extracted and pushed down rather than collapsing to the universe interval.
+        """
+        ts = [datetime(2024, 1, 1) + timedelta(days=i) for i in range(90)]
+        df = pl.DataFrame({"ts": ts, "val": list(range(90))})
+        tracker = PredicateTracker(df)
+
+        lf = multi_source(
+            sources={"data": (tracker.lazy_frame, {"ts": FilterSpec()})},
+            combine=lambda s: s["data"],
+        )
+
+        start, end = date(2024, 1, 1), date(2024, 1, 31)
+        result = lf.filter((pl.col("ts").dt.date() >= start) & (pl.col("ts").dt.date() <= end)).collect()
+
+        # Range was extracted, so the source received a constraining predicate (not a silent full scan).
+        assert tracker.last_predicate is not None
+        assert result.height == 31
+        assert result["ts"].max() == datetime(2024, 1, 31)
+
     def test_simple_equality_filter(self):
         """Equality filter propagates correctly."""
         left = pl.LazyFrame({"id": ["A", "B", "C"], "val": [1, 2, 3]})