Allow strict casts of literals only in polars lazyframe pushdown (#348)

Author: evertlammerts

duckdb/polars_io.py

@@ -170,8 +170,17 @@ def _pl_tree_to_sql(tree: _ExpressionTree) -> str:
     if node_type == "Cast":
         cast_tree = tree[node_type]
         assert isinstance(cast_tree, dict), f"A {node_type} should be a dict but got {type(cast_tree)}"
-        if cast_tree.get("options") not in ("NonStrict", "Strict"):
-            msg = f"Only NonStrict/Strict casts can be safely unwrapped, got {cast_tree.get('options')!r}"
+        options = cast_tree.get("options")
+        if options == "Strict":
+            # Strict casts on literals (e.g. pl.lit(1, dtype=pl.Int8)) are safe to unwrap —
+            # the value is known at expression creation time. Strict casts on columns
+            # (e.g. pl.col("a").cast(pl.Int64)) are semantically meaningful and must not be dropped.
+            cast_expr = cast_tree.get("expr", {})
+            if not isinstance(cast_expr, dict) or "Literal" not in cast_expr:
+                msg = "Strict cast on non-literal expression cannot be pushed down"
+                raise NotImplementedError(msg)
+        elif options != "NonStrict":
+            msg = f"Only NonStrict/Strict casts can be safely unwrapped, got {options!r}"
             raise NotImplementedError(msg)
         cast_expr = cast_tree["expr"]
         assert isinstance(cast_expr, dict), f"A {node_type} should be a dict but got {type(cast_expr)}"

tests/fast/arrow/test_polars.py

@@ -12,6 +12,7 @@
 
 from duckdb.polars_io import _pl_tree_to_sql, _predicate_to_expression  # noqa: E402
 
+pl_pre_1_35_0 = parse_version(pl.__version__) < parse_version("1.35.0")
 pl_pre_1_36_0 = parse_version(pl.__version__) < parse_version("1.36.0")
 
 
@@ -437,7 +438,7 @@ def test_polars_lazy_pushdown_timestamp(self, duckdb_cursor):
             lazy_df.filter((pl.col("a") == ts_2020) | (pl.col("b") == ts_2008)).select(pl.len()).collect().item() == 2
         )
 
-    @pytest.mark.skipif(pl_pre_1_36_0, reason="Polars < 1.36.0 expressions on dates produce casts in predicates")
+    @pytest.mark.skipif(pl_pre_1_35_0, reason="Polars < 1.36.0 expressions on dates produce casts in predicates")
     def test_polars_predicate_to_expression_post_1_36_0(self):
         ts_2008 = datetime.datetime(2008, 1, 1, 0, 0, 1)
         ts_2010 = datetime.datetime(2010, 1, 1, 10, 0, 1)
@@ -454,7 +455,7 @@ def test_polars_predicate_to_expression_post_1_36_0(self):
         valid_filter((pl.col("a") == ts_2020) & (pl.col("b") == ts_2010) & (pl.col("c") == ts_2020))
         valid_filter((pl.col("a") == ts_2020) | (pl.col("b") == ts_2008))
 
-    @pytest.mark.skipif(not pl_pre_1_36_0, reason="Polars >= 1.36.0 expressions on dates don't produce casts")
+    @pytest.mark.skipif(not pl_pre_1_35_0, reason="Polars >= 1.36.0 expressions on dates don't produce casts")
     def test_polars_predicate_to_expression_pre_1_36_0(self):
         ts_2008 = datetime.datetime(2008, 1, 1, 0, 0, 1)
         ts_2010 = datetime.datetime(2010, 1, 1, 10, 0, 1)