Drop dynamic predicates from logical plan

python/cudf_polars/cudf_polars/dsl/translate.py

@@ -286,15 +286,42 @@ def __exit__(self, *args: Any) -> None:
         self.translator._expr_context = self._prev
 
 
-def _is_dynamic_pred(visitor: Any, expr_ir: Any) -> bool:
+def _drop_dyn_pred_hints(
+    translator: Translator, n: int, schema: Schema
+) -> expr.Expr | None:
     try:
-        visitor.view_expression(expr_ir.node)
+        node = translator.visitor.view_expression(n)
     except Exception as e:
         if str(e) == "dynamic_pred":
-            return True
-        raise  # pragma: no cover
-    else:
-        return False
+            return None
+        raise
+    if isinstance(node, plrs._expr_nodes.BinaryExpr) and node.op in (
+        plrs._expr_nodes.Operator.And,
+        plrs._expr_nodes.Operator.LogicalAnd,
+    ):
+        left = _drop_dyn_pred_hints(translator, node.left, schema)
+        right = _drop_dyn_pred_hints(translator, node.right, schema)
+        if left is None:
+            return right
+        if right is None:
+            return left
+        return expr.BinOp(
+            DataType(translator.visitor.get_dtype(n)),
+            expr.BinOp._MAPPING[node.op],
+            left,
+            right,
+        )
+    return translator.translate_expr(n=n, schema=schema)
+
+
+def translate_predicate(
+    translator: Translator, *, n: Any, schema: Schema
+) -> expr.NamedExpr | None:
+    """Translate predicate, dropping Polars' dynamic predicate hints."""
+    mask = _drop_dyn_pred_hints(translator, n.node, schema)
+    if mask is None:
+        return None
+    return expr.NamedExpr(n.output_name, mask)
 
 
 class set_internal_name_gen(AbstractContextManager[None]):
@@ -397,8 +424,7 @@ def _(node: plrs._ir_nodes.Scan, translator: Translator, schema: Schema) -> ir.I
         (
             None
             if node.predicate is None
-            or _is_dynamic_pred(translator.visitor, node.predicate)
-            else translate_named_expr(translator, n=node.predicate, schema=schema)
+            else translate_predicate(translator, n=node.predicate, schema=schema)
         ),
         parquet_options,
     )
@@ -615,9 +641,9 @@ def _(node: plrs._ir_nodes.Slice, translator: Translator, schema: Schema) -> ir.
 def _(node: plrs._ir_nodes.Filter, translator: Translator, schema: Schema) -> ir.IR:
     with set_node(translator.visitor, node.input):
         inp = translator.translate_ir(n=None)
-        if _is_dynamic_pred(translator.visitor, node.predicate):
+        mask = translate_predicate(translator, n=node.predicate, schema=inp.schema)
+        if mask is None:
             return inp
-        mask = translate_named_expr(translator, n=node.predicate, schema=inp.schema)
     return ir.Filter(schema, mask, inp)
 
 

python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py

@@ -1164,9 +1164,9 @@ def run_polars_in_memory(
         **run_config.streaming_options.to_engine_options(),
         "parquet_options": parquet_options,
     }
+    engine_options.setdefault("raise_on_fail", True)
     engine = pl.GPUEngine(
         executor="in-memory",
-        raise_on_fail=True,
         **engine_options,
     )
     records, plans, validation_failures, query_failures = _run_query_loop(
@@ -1202,6 +1202,7 @@ def run_polars_spmd(
         **run_config.streaming_options.to_engine_options(),
         "parquet_options": parquet_options,
     }
+    engine_options.setdefault("raise_on_fail", True)
     with SPMDEngine(
         rapidsmpf_options=run_config.streaming_options.to_rapidsmpf_options(),
         executor_options=executor_options,
@@ -1259,6 +1260,7 @@ def run_polars_ray(
         **run_config.streaming_options.to_engine_options(),
         "parquet_options": parquet_options,
     }
+    engine_options.setdefault("raise_on_fail", True)
     ray_init_options: dict[str, Any] = {}
     if run_config.connect is not None:
         ray_init_options["address"] = run_config.connect
@@ -1308,6 +1310,7 @@ def run_polars_dask(
         **run_config.streaming_options.to_engine_options(),
         "parquet_options": parquet_options,
     }
+    engine_options.setdefault("raise_on_fail", True)
     dask_client = None
     if run_config.connect is not None:
         if Path(run_config.connect).is_file():

python/cudf_polars/tests/test_filter.py

@@ -5,7 +5,11 @@
 import pytest
 
 import polars as pl
+from polars import polars as plrs
 
+from cudf_polars import Translator
+from cudf_polars.dsl import expr, ir
+from cudf_polars.dsl.traversal import traversal
 from cudf_polars.testing.asserts import assert_gpu_result_equal
 
 
@@ -28,3 +32,38 @@ def test_filter(engine: pl.GPUEngine, expr, predicate_pushdown):
             "optimizations": pl.QueryOptFlags(predicate_pushdown=predicate_pushdown)
         },
     )
+
+
+def test_filter_drops_dynamic_predicate_hint():
+    class DynamicPredHint(Exception):
+        pass
+
+    class DynamicPredVisitor:
+        def __init__(self, inner, op):
+            self.inner = inner
+            self.op = op
+
+        def __getattr__(self, name):
+            return getattr(self.inner, name)
+
+        def view_expression(self, n):
+            node = self.inner.view_expression(n)
+            if isinstance(node, plrs._expr_nodes.BinaryExpr) and node.op == self.op:
+                raise DynamicPredHint("dynamic_pred")
+            return node
+
+    ldf = pl.LazyFrame(
+        {"a": [1, 2, 3, 4, 5], "b": [5, 4, 3, 2, 1], "c": [1, 1, 3, 3, 5]}
+    )
+    query = ldf.filter(
+        (pl.col("b") < 5) & (pl.col("a") > 1) & (pl.col("c") == 3) & (pl.col("b") < 3)
+    )
+
+    visitor = DynamicPredVisitor(query._ldf.visit(), plrs._expr_nodes.Operator.Lt)
+    translator = Translator(visitor, pl.GPUEngine())
+    node = translator.translate_ir()
+
+    assert translator.errors == []
+    assert isinstance(node, ir.Filter)
+    columns = {n.name for n in traversal([node.mask.value]) if isinstance(n, expr.Col)}
+    assert columns == {"a", "c"}