diff --git a/python/cudf_polars/cudf_polars/dsl/translate.py b/python/cudf_polars/cudf_polars/dsl/translate.py index d1b86d40ee0a..0e882ed0f448 100644 --- a/python/cudf_polars/cudf_polars/dsl/translate.py +++ b/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 + return None raise # pragma: no cover - else: - return False + 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) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 21b3787ae74c..8b83b31ed107 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -1210,9 +1210,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( @@ -1248,6 +1248,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, @@ -1305,6 +1306,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 @@ -1354,6 +1356,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(): diff --git a/python/cudf_polars/tests/test_filter.py b/python/cudf_polars/tests/test_filter.py index 20dc656470f6..71781454e70d 100644 --- a/python/cudf_polars/tests/test_filter.py +++ b/python/cudf_polars/tests/test_filter.py @@ -28,3 +28,15 @@ def test_filter(engine: pl.GPUEngine, expr, predicate_pushdown): "optimizations": pl.QueryOptFlags(predicate_pushdown=predicate_pushdown) }, ) + + +def test_filter_drops_dynamic_predicate_hint(engine: pl.GPUEngine): + ldf = pl.LazyFrame( + {"a": [1, 2, 3, 4, 5], "b": [5, 4, 3, 2, 1], "c": [1, 1, 3, 3, 5]} + ) + # sort("b").head(3) causes Polars to inject a dynamic predicate hint into + # the filter below: FILTER (a > 1) & (c == 3) & col("b").dynamic_predicate() + # This test ensures we drop these dynamic predicate hints from the filter + # before executing on the GPU. + q = ldf.filter((pl.col("a") > 1) & (pl.col("c") == 3)).sort("b").head(3) + assert_gpu_result_equal(q, engine=engine)