diff --git a/python/cudf_polars/cudf_polars/streaming/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py index 8287980277e0..f0680586d8fa 100644 --- a/python/cudf_polars/cudf_polars/streaming/explain.py +++ b/python/cudf_polars/cudf_polars/streaming/explain.py @@ -14,10 +14,15 @@ from itertools import groupby from typing import TYPE_CHECKING, Any, Self, TypeAlias -import cudf_polars.dsl.expressions.binaryop -import cudf_polars.dsl.expressions.literal -from cudf_polars.dsl.expressions.base import NamedExpr +import pylibcudf as plc + +from cudf_polars.dsl.expressions.base import Col, ColRef, Expr, NamedExpr +from cudf_polars.dsl.expressions.binaryop import BinOp +from cudf_polars.dsl.expressions.literal import Literal +from cudf_polars.dsl.expressions.ternary import Ternary +from cudf_polars.dsl.expressions.unary import Cast, UnaryFunction from cudf_polars.dsl.ir import ( + ConditionalJoin, Filter, GroupBy, HStack, @@ -257,6 +262,59 @@ def _(ir: Join, *, offset: str = "") -> str: return _repr_header(offset, f"JOIN {ir.options[0]} {left_on} {right_on}", ir.schema) +_BinaryOperator = plc.binaryop.BinaryOperator +_BINOP_SYMBOLS: dict[_BinaryOperator, str] = { + _BinaryOperator.EQUAL: "==", + _BinaryOperator.NOT_EQUAL: "!=", + _BinaryOperator.LESS: "<", + _BinaryOperator.LESS_EQUAL: "<=", + _BinaryOperator.GREATER: ">", + _BinaryOperator.GREATER_EQUAL: ">=", + _BinaryOperator.LOGICAL_AND: "&", + _BinaryOperator.NULL_LOGICAL_AND: "&", + _BinaryOperator.LOGICAL_OR: "|", + _BinaryOperator.NULL_LOGICAL_OR: "|", +} + + +def _predicate_to_str(expr: Expr) -> str: + match expr: + case Col(name=name): + return name + case ColRef(): + col = expr.children[0] + assert isinstance(col, Col) + return col.name + case Literal(value=value): + return repr(value) + case Cast(): + return _predicate_to_str(expr.children[0]) + case BinOp(op=op): + left, right = expr.children + sym = _BINOP_SYMBOLS.get(op, op.name) + return f"({_predicate_to_str(left)} {sym} {_predicate_to_str(right)})" + case UnaryFunction(name=name): + (child,) = expr.children + return f"{name}({_predicate_to_str(child)})" + case Ternary(): + when, then, otherwise = expr.children + return f"when({_predicate_to_str(when)}).then({_predicate_to_str(then)}).otherwise({_predicate_to_str(otherwise)})" + case _: + return type(expr).__name__ + + +@_repr_ir.register +def _(ir: ConditionalJoin, *, offset: str = "") -> str: + pred = _predicate_to_str(ir.predicate) + return _repr_header(offset, f"CONDITIONALJOIN {pred}", ir.schema) + + +@_repr_ir.register +def _(ir: Filter, *, offset: str = "") -> str: + pred = _predicate_to_str(ir.mask.value) + return _repr_header(offset, f"FILTER {pred}", ir.schema) + + @_repr_ir.register def _(ir: Sort, *, offset: str = "") -> str: by = tuple(ne.name for ne in ir.by) @@ -266,6 +324,8 @@ def _(ir: Sort, *, offset: str = "") -> str: @_repr_ir.register def _(ir: Scan, *, offset: str = "") -> str: label = f"SCAN {ir.typ.upper()}" + if ir.predicate is not None: + label += f" {_predicate_to_str(ir.predicate.value)}" return _repr_header(offset, label, ir.schema) @@ -344,11 +404,11 @@ def _serialize_expr(expr: Expr | NamedExpr) -> dict[str, Serializable]: match expr: case NamedExpr(name=name, value=value): return {"type": "NamedExpr", "name": name, "value": _serialize_expr(value)} - case cudf_polars.dsl.expressions.base.Col(name=name): + case Col(name=name): return {"type": "Col", "name": name} - case cudf_polars.dsl.expressions.literal.Literal(value=value): + case Literal(value=value): return {"type": "Literal", "value": _serialize_literal(value)} - case cudf_polars.dsl.expressions.binaryop.BinOp(): + case BinOp(): return { "op": expr.op.name, "left": _serialize_expr(expr.children[0]), diff --git a/python/cudf_polars/tests/streaming/test_explain.py b/python/cudf_polars/tests/streaming/test_explain.py index 5f4984d53123..41e0177dad10 100644 --- a/python/cudf_polars/tests/streaming/test_explain.py +++ b/python/cudf_polars/tests/streaming/test_explain.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -13,9 +13,15 @@ import polars as pl +import pylibcudf as plc + +from cudf_polars.containers import DataType +from cudf_polars.dsl.expressions.base import Col +from cudf_polars.dsl.expressions.binaryop import BinOp from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.explain import ( _fmt_row_count, + _predicate_to_str, explain_query, serialize_query, ) @@ -658,6 +664,44 @@ def test_hstack_properties(): assert node.properties == {"columns": ["a", "b"]} +def test_predicate_to_str_col(): + col = Col(DataType(pl.Float64()), "c_acctbal") + assert _predicate_to_str(col) == "c_acctbal" + + +def test_predicate_to_str_binop(): + float_dtype = DataType(pl.Float64()) + bool_dtype = DataType(pl.Boolean()) + left = Col(float_dtype, "c_acctbal") + right = Col(float_dtype, "avg_acctbal") + expr = BinOp(bool_dtype, plc.binaryop.BinaryOperator.GREATER, left, right) + assert _predicate_to_str(expr) == "(c_acctbal > avg_acctbal)" + + +def test_predicate_to_str_nested_binop(): + float_dtype = DataType(pl.Float64()) + bool_dtype = DataType(pl.Boolean()) + a = Col(float_dtype, "a") + b = Col(float_dtype, "b") + c = Col(float_dtype, "c") + ab = BinOp(bool_dtype, plc.binaryop.BinaryOperator.GREATER, a, b) + abc = BinOp(bool_dtype, plc.binaryop.BinaryOperator.LOGICAL_AND, ab, c) + assert _predicate_to_str(abc) == "((a > b) & c)" + + +def test_explain_conditional_join_shows_predicate(): + customers = pl.LazyFrame({"c_val": [4.0, 5.0, 6.0]}) + avg = pl.LazyFrame({"avg_val": [3.5]}) + q = customers.join_where(avg, pl.col("c_val") > pl.col("avg_val")) + + engine = pl.GPUEngine(executor="streaming", raise_on_fail=True) + with pytest.warns(UserWarning, match="ConditionalJoin not supported"): + plan = explain_query(q, engine, physical=True) + + assert "CONDITIONALJOIN" in plan + assert "c_val > avg_val" in plan + + def test_explain_physical_plan(tmp_path, df): make_partitioned_source(df, tmp_path, fmt="parquet", n_files=5)