diff --git a/python/cudf_polars/cudf_polars/dsl/nodebase.py b/python/cudf_polars/cudf_polars/dsl/nodebase.py index 23aaf4138a08..6453c798f7fc 100644 --- a/python/cudf_polars/cudf_polars/dsl/nodebase.py +++ b/python/cudf_polars/cudf_polars/dsl/nodebase.py @@ -5,6 +5,7 @@ from __future__ import annotations +import hashlib from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar if TYPE_CHECKING: @@ -33,8 +34,9 @@ class Node(Generic[T]): *children).`` """ - __slots__ = ("_hash_value", "_repr_value", "children") + __slots__ = ("_hash_value", "_repr_value", "_stable_hash_value", "children") _hash_value: int + _stable_hash_value: int _repr_value: str children: tuple[T, ...] _non_child: ClassVar[tuple[str, ...]] = () @@ -81,6 +83,30 @@ def get_hashable(self) -> Hashable: """ return (type(self), self._ctor_arguments(self.children)) + def get_stable_id(self) -> int: + """ + Compute a stable identifier for Node. + + Uses MD5 hash of the node's hashable representation for determinism + across process boundaries (Python's hash() uses PYTHONHASHSEED). + + Parameters + ---------- + ir_node + The IR node. + + Returns + ------- + int + A stable 32-bit identifier for this node. + """ + try: + return self._stable_hash_value + except AttributeError: + content = repr(self.get_hashable()).encode("utf-8") + self._stable_hash_value = int(hashlib.md5(content).hexdigest()[:8], 16) + return self._stable_hash_value + def __hash__(self) -> int: """ Hash of an expression with caching. diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py index 7e0da8c32d17..e30a141f6891 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py @@ -59,6 +59,8 @@ if TYPE_CHECKING: from collections.abc import Callable, Sequence + from cudf_polars.experimental.explain import SerializablePlan + try: import structlog @@ -237,6 +239,7 @@ class RunConfig: default_factory=PackageVersions.collect ) records: dict[int, list[Record]] = dataclasses.field(default_factory=dict) + plans: dict[int, SerializablePlan] = dataclasses.field(default_factory=dict) dataset_path: Path scale_factor: int | float shuffle: Literal["rapidsmpf", "tasks"] | None = None @@ -507,28 +510,37 @@ def print_query_plan( args: argparse.Namespace, run_config: RunConfig, engine: None | pl.GPUEngine = None, -) -> None: + *, + print_plans: bool = True, +) -> tuple[str | None, str | None]: """Print the query plan.""" + logical_plan = plan = None if run_config.executor == "cpu": if args.explain_logical: - print(f"\nQuery {q_id} - Logical plan\n") - print(q.explain()) + logical_plan = q.explain() if args.explain: - print(f"\nQuery {q_id} - Physical plan\n") - print(q.show_graph(engine="streaming", plan_stage="physical")) + plan = q.show_graph(engine="streaming", plan_stage="physical") elif CUDF_POLARS_AVAILABLE: assert isinstance(engine, pl.GPUEngine) if args.explain_logical: - print(f"\nQuery {q_id} - Logical plan\n") - print(explain_query(q, engine, physical=False)) + logical_plan = explain_query(q, engine, physical=False) if args.explain and run_config.executor == "streaming": - print(f"\nQuery {q_id} - Physical plan\n") - print(explain_query(q, engine)) + plan = explain_query(q, engine) else: raise RuntimeError( "Cannot provide the logical or physical plan because cudf_polars is not installed." ) + if print_plans: + if logical_plan: + print(f"\nQuery {q_id} - Logical plan\n") + print(logical_plan) + if plan: + print(f"\nQuery {q_id} - Physical plan\n") + print(plan) + + return logical_plan, plan + def initialize_dask_cluster(run_config: RunConfig, args: argparse.Namespace): # type: ignore[no-untyped-def] """ @@ -947,6 +959,12 @@ def parse_args( help="Print an outline of the logical plan", default=False, ) + parser.add_argument( + "--print-plans", + action=argparse.BooleanOptionalAction, + help="Print the query plans", + default=True, + ) parser.add_argument( "--validate", action=argparse.BooleanOptionalAction, @@ -1052,6 +1070,7 @@ def run_polars( run_config = dataclasses.replace(run_config, n_workers=actual_n_workers) records: defaultdict[int, list[Record]] = defaultdict(list) + plans: dict[int, SerializablePlan] = {} engine: pl.GPUEngine | None = None if run_config.executor != "cpu": @@ -1079,7 +1098,13 @@ def run_polars( except AttributeError as err: raise NotImplementedError(f"Query {q_id} not implemented.") from err - print_query_plan(q_id, q, args, run_config, engine) + print_query_plan( + q_id, q, args, run_config, engine, print_plans=args.print_plans + ) + if (args.explain or args.explain_logical) and engine is not None: + from cudf_polars.experimental.explain import serialize_query + + plans[q_id] = serialize_query(q, engine) records[q_id] = [] for i in range(args.iterations): @@ -1140,7 +1165,7 @@ def run_polars( ) records[q_id].append(record) - run_config = dataclasses.replace(run_config, records=dict(records)) + run_config = dataclasses.replace(run_config, records=dict(records), plans=plans) # consolidate logs if _HAS_STRUCTLOG and run_config.collect_traces: diff --git a/python/cudf_polars/cudf_polars/experimental/explain.py b/python/cudf_polars/cudf_polars/experimental/explain.py index 17cba1d17cb4..ecd72ef46c04 100644 --- a/python/cudf_polars/cudf_polars/experimental/explain.py +++ b/python/cudf_polars/cudf_polars/experimental/explain.py @@ -1,22 +1,28 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 """Explain logical and physical plans.""" from __future__ import annotations +import dataclasses import functools +from collections.abc import Mapping, Sequence from itertools import groupby -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Self, TypeAlias +import cudf_polars.dsl.expressions.binaryop +import cudf_polars.dsl.expressions.literal from cudf_polars.dsl.ir import ( + Filter, GroupBy, + HStack, Join, Scan, + Select, Sort, ) from cudf_polars.dsl.translate import Translator +from cudf_polars.dsl.traversal import traversal from cudf_polars.experimental.base import ColumnStat from cudf_polars.experimental.parallel import lower_ir_graph from cudf_polars.experimental.statistics import ( @@ -29,10 +35,22 @@ import polars as pl + from cudf_polars.dsl.expressions.base import Expr from cudf_polars.dsl.ir import IR from cudf_polars.experimental.base import PartitionInfo, StatsCollector +Serializable: TypeAlias = ( + str + | int + | float + | bool + | Sequence["Serializable"] + | Mapping[str, "Serializable"] + | None +) + + def explain_query( q: pl.LazyFrame, engine: pl.GPUEngine, @@ -81,6 +99,76 @@ def explain_query( return _repr_ir_tree(ir) +def serialize_query( + q: pl.LazyFrame, + engine: pl.GPUEngine, + *, + physical: bool = True, +) -> SerializablePlan: + """ + Return a structured, serializable representation of the IR plan. + + Parameters + ---------- + q : pl.LazyFrame + The LazyFrame to serialize. + engine : pl.GPUEngine + The configured GPU engine to use. + physical : bool, default True + If True, serialize the physical (lowered) plan with partition info. + If False, serialize the logical (pre-lowering) plan. + + Returns + ------- + plan + A structured representation of the query plan that can be + serialized to JSON. + + Examples + -------- + >>> import polars as pl + >>> import json + >>> import dataclasses + >>> q = pl.LazyFrame({"a": [1, 2, 3]}).select(pl.col("a") * 2) + >>> engine = pl.GPUEngine(executor="streaming") + >>> plan = serialize_query(q, engine, physical=False) + >>> print(json.dumps(dataclasses.asdict(plan), indent=2)) + { + "roots": [ + "1739020873" + ], + "nodes": { + "1739020873": { + "id": "1739020873", + "children": [ + "2653195019" + ], + "schema": { + "a": "INT64" + }, + "properties": { + "columns": [ + "a" + ] + }, + "type": "Select" + }, + "2653195019": { + "id": "2653195019", + "children": [], + "schema": { + "a": "INT64" + }, + "properties": {}, + "type": "DataFrameScan" + } + }, + "partition_info": null + } + """ + return SerializablePlan.from_query(q, engine, lowered=physical) + + def _fmt_row_count(value: int | None) -> str: """Format a row count as a readable string.""" if value is None: @@ -168,3 +256,242 @@ def _(ir: Sort, *, offset: str = "") -> str: def _(ir: Scan, *, offset: str = "") -> str: label = f"SCAN {ir.typ.upper()}" return _repr_header(offset, label, ir.schema) + + +# -------------------------------------------------------------------------- +# Property serialization for structured query plan export +# -------------------------------------------------------------------------- + + +@functools.singledispatch +def _serialize_properties(ir: IR) -> dict[str, Serializable]: + """Extract serializable properties from an IR node.""" + return {} + + +@_serialize_properties.register +def _(ir: Scan) -> dict[str, Serializable]: + # for polars<1.31, paths is a list[Path] + # for polars>=1.31, paths is a list[str] + return { + "typ": ir.typ, + "paths": [str(path) for path in ir.paths], + } + + +@_serialize_properties.register +def _(ir: Join) -> dict[str, Serializable]: + return { + "how": ir.options[0], + "left_on": [ne.name for ne in ir.left_on], + "right_on": [ne.name for ne in ir.right_on], + } + + +@_serialize_properties.register +def _(ir: GroupBy) -> dict[str, Serializable]: + return { + "keys": [ne.name for ne in ir.keys], + } + + +@_serialize_properties.register +def _(ir: Sort) -> dict[str, Serializable]: + return { + "by": [ne.name for ne in ir.by], + "order": [o.name for o in ir.order], + } + + +def _serialize_expr(expr: Expr) -> dict[str, Serializable]: + match expr: + case cudf_polars.dsl.expressions.base.Col(name=name): + return {"type": "Col", "name": name} + case cudf_polars.dsl.expressions.literal.Literal(value=value): + return {"type": "Literal", "value": value} + case cudf_polars.dsl.expressions.binaryop.BinOp(): + return { + "op": expr.op.name, + "left": _serialize_expr(expr.children[0]), + "right": _serialize_expr(expr.children[1]), + } + case _: # pragma: no cover + return {"type": type(expr).__name__} + + +@_serialize_properties.register +def _(ir: Filter) -> dict[str, Serializable]: + value = ir.mask.value + properties = _serialize_expr(value) + properties["predicate"] = ir.mask.name + + return properties + + +@_serialize_properties.register +def _(ir: Select) -> dict[str, Serializable]: + return { + "columns": [ne.name for ne in ir.exprs], + } + + +@_serialize_properties.register +def _(ir: HStack) -> dict[str, Serializable]: + return { + "columns": [ne.name for ne in ir.columns], + } + + +@dataclasses.dataclass +class SerializableIRNode: + """ + A node in the plan. + + This node is *serializable* and cannot be executed like a + cudf_polars.dsl.ir.IR node. + """ + + id: str + children: list[str] + schema: dict[str, Serializable] + properties: dict[str, Serializable] + type: str + + @classmethod + def from_ir(cls, ir: IR) -> Self: + """Build a Node from an IR Node.""" + return cls( + id=str(ir.get_stable_id()), + children=[str(child.get_stable_id()) for child in ir.children], + schema={k: v.id().name for k, v in ir.schema.items()}, + properties=_serialize_properties(ir), + type=type(ir).__name__, + ) + + +@dataclasses.dataclass +class SerializablePartitionInfo: + """Serializable information about a partition.""" + + count: int + partitioned_on: tuple[Serializable, ...] + + +@dataclasses.dataclass +class SerializablePlan: + """ + A serializable representation of a query plan. + + Parameters + ---------- + roots + The IDs of the root nodes of the plan. + nodes + A mapping from node ID to node details. + partition_info + Information about the partitions of the plan. + + Notes + ----- + All integers node IDs are stored as strings to make round-tripping + to JSON easier. Node IDs will appear in + + - ``roots`` + - the keys of ``nodes`` + - the ``children`` of each node in ``nodes`` + - the keys in ``partition_info`` + + You can safely rely on every key being present in ``nodes``. + + See Also + -------- + serialize_query + A function that builds a serializable plan from a LazyFrame query. + """ + + roots: list[str] + nodes: dict[str, SerializableIRNode] + partition_info: dict[str, SerializablePartitionInfo] | None = None + + @classmethod + def from_ir( + cls, ir: IR, *, config_options: ConfigOptions, lowered: bool = False + ) -> Self: + """ + Construct a serializable plan from an IR node. + + Parameters + ---------- + ir + The IR node to construct the serializable plan from. + config_options + The configuration options. + lowered + If True, lower the IR to the physical plan and include partition info. + + Returns + ------- + plan + A serializable representation of the query plan. + """ + partition_info_dict: dict[str, SerializablePartitionInfo] | None = None + if lowered: + if ( + config_options.executor.name == "streaming" + and config_options.executor.runtime == "rapidsmpf" + ): # pragma: no cover; rapidsmpf runtime not tested in CI yet + from cudf_polars.experimental.rapidsmpf.core import ( + lower_ir_graph as rapidsmpf_lower_ir_graph, + ) + + ir, partition_info_d, _ = rapidsmpf_lower_ir_graph(ir, config_options) + else: + ir, partition_info_d, _ = lower_ir_graph(ir, config_options) + partition_info_dict = {} + + nodes: dict[str, SerializableIRNode] = {} + for ir_node in traversal([ir]): + stable_id = str(ir_node.get_stable_id()) + nodes[stable_id] = SerializableIRNode.from_ir(ir_node) + if partition_info_dict is not None: + partition_info_dict[stable_id] = SerializablePartitionInfo( + count=partition_info_d[ir_node].count, + partitioned_on=tuple( + expr.name for expr in partition_info_d[ir_node].partitioned_on + ), + ) + + return cls( + roots=[str(ir.get_stable_id())], + nodes=nodes, + partition_info=partition_info_dict, + ) + + @classmethod + def from_query( + cls, + q: pl.LazyFrame, + engine: pl.GPUEngine, + *, + lowered: bool = False, + ) -> Self: + """ + Build a serializable plan from a LazyFrame query. + + Parameters + ---------- + q + The LazyFrame to serialize. + engine + The GPU engine to use. If None, uses default streaming executor. + lowered + If True, lower the IR to the physical plan and include partition info. + + Returns + ------- + plan + A serializable representation of the query plan. + """ + config_options = ConfigOptions.from_polars_engine(engine) + ir = Translator(q._ldf.visit(), engine).translate_ir() + return cls.from_ir(ir, config_options=config_options, lowered=lowered) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py index cdc8bd04d952..e405a383456a 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py @@ -98,7 +98,7 @@ def evaluate_logical_plan( ir, partition_info, stats = lower_ir_graph(ir, config_options) # Log the query plan structure for tracing (no-op if tracing disabled) - log_query_plan(ir) + log_query_plan(ir, config_options) # Reserve shuffle IDs for the entire pipeline execution with ReserveOpIDs(ir, config_options) as collective_id_map: diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py index 31ee9c113bb2..ec24bfab3b7a 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py @@ -4,37 +4,17 @@ from __future__ import annotations -import hashlib +import dataclasses from typing import TYPE_CHECKING from cudf_polars.dsl.tracing import LOG_TRACES, Scope -from cudf_polars.dsl.traversal import traversal +from cudf_polars.experimental.explain import SerializablePlan if TYPE_CHECKING: import pylibcudf as plc from cudf_polars.dsl.ir import IR - - -def _stable_ir_id(ir_node: IR) -> int: - """ - Compute a stable identifier for an IR node. - - Uses MD5 hash of the node's hashable representation for determinism - across process boundaries (Python's hash() uses PYTHONHASHSEED). - - Parameters - ---------- - ir_node - The IR node. - - Returns - ------- - int - A stable 32-bit identifier for this node. - """ - content = repr(ir_node.get_hashable()).encode("utf-8") - return int(hashlib.md5(content).hexdigest()[:8], 16) + from cudf_polars.utils.config import ConfigOptions class ActorTracer: @@ -100,7 +80,7 @@ def set_duplicated(self, *, duplicated: bool = True) -> None: self.duplicated = duplicated -def log_query_plan(ir: IR) -> None: +def log_query_plan(ir: IR, config_options: ConfigOptions) -> None: """ Log the IR tree structure as a structlog event. @@ -112,6 +92,8 @@ def log_query_plan(ir: IR) -> None: ---------- ir The root IR node of the lowered query plan. + config_options + The GPU engine configuration options. Notes ----- @@ -122,14 +104,8 @@ def log_query_plan(ir: IR) -> None: import structlog - nodes = [ - { - "ir_id": _stable_ir_id(node), - "ir_type": type(node).__name__, - "children_ir_ids": [_stable_ir_id(c) for c in node.children], - } - for node in traversal([ir]) - ] + dag = SerializablePlan.from_ir(ir, config_options=config_options) + raw = dataclasses.asdict(dag) log = structlog.get_logger() - log.info("Query Plan", scope=Scope.PLAN.value, nodes=nodes) + log.info("Query Plan", scope=Scope.PLAN.value, plan=raw) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py index 14516dae15fa..8463226ae02e 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py @@ -82,10 +82,9 @@ async def shutdown_on_error( if LOG_TRACES and trace_ir is not None: from cudf_polars.experimental.rapidsmpf.tracing import ( ActorTracer, - _stable_ir_id, ) - ir_id = _stable_ir_id(trace_ir) + ir_id = trace_ir.get_stable_id() ir_type = type(trace_ir).__name__ tracer = ActorTracer(ir_id, ir_type) structlog.contextvars.bind_contextvars(actor_ir_id=ir_id, actor_ir_type=ir_type) diff --git a/python/cudf_polars/docs/overview.md b/python/cudf_polars/docs/overview.md index 4a1351dba8d2..3914c68c3319 100644 --- a/python/cudf_polars/docs/overview.md +++ b/python/cudf_polars/docs/overview.md @@ -796,3 +796,52 @@ These provide a higher-level grouping over the lower-level libcudf calls (e.g. Finally, if using [rapidsmpf](https://docs.rapids.ai/api/rapidsmpf/nightly/) for shuffling, the methods inserting and extracting partitions to shuffle are annotated with nvtx ranges. + +# Query Plans + +The module `cudf_polars.experimental.explain` contains functions for dumping +the query for a given `LazyFrame`. + + +## Structured Output + +`cudf_polars.experimental.explain.serialize_query` can be used to output +the query plan in a structured format. + +```python +>>> import dataclasses +>>> import polars as pl +>>> from cudf_polars.experimental.explain import serialize_query +>>> q = pl.LazyFrame({"a": ['a', 'b', 'a'], "b": [1, 2, 3]}).group_by("a").agg(pl.len()) +>>> dataclasses.asdict(serialize_query(q, engine=pl.GPUEngine())) +{'roots': ['526964741'], + 'nodes': {'526964741': {'id': '526964741', + 'children': ['1694929589'], + 'schema': {'a': 'STRING', 'len': 'UINT32'}, + 'properties': {'columns': ['a', 'len']}, + 'type': 'Select'}, + '1694929589': {'id': '1694929589', + 'children': ['2632275007'], + 'schema': {'a': 'STRING', '___0': 'UINT32'}, + 'properties': {'keys': ['a']}, + 'type': 'GroupBy'}, + '2632275007': {'id': '2632275007', + 'children': [], + 'schema': {'a': 'STRING'}, + 'properties': {}, + 'type': 'DataFrameScan'}}, + 'partition_info': {'526964741': {'count': 1, 'partitioned_on': ()}, + '1694929589': {'count': 1, 'partitioned_on': ()}, + '2632275007': {'count': 1, 'partitioned_on': ()}}} +``` + +The structured schema has three top-level fields: + +1. `roots`: the integer ID for the "root" (final) nodes in the query plan +2. `partition_info`: partitioning information at each stage of the query +3. `nodes`: A mapping from integer node id to node details. Each node ID + that appears in the output will be present in this mapping. + Inspect `children` to understand which nodes this node depends on. + +Note that all integers are stored as strings to make round-tripping +to JSON easier. diff --git a/python/cudf_polars/tests/experimental/test_explain.py b/python/cudf_polars/tests/experimental/test_explain.py index 5169c793b15b..61d2cb76f402 100644 --- a/python/cudf_polars/tests/experimental/test_explain.py +++ b/python/cudf_polars/tests/experimental/test_explain.py @@ -3,13 +3,20 @@ from __future__ import annotations +import dataclasses +import json import re +from typing import TYPE_CHECKING import pytest import polars as pl -from cudf_polars.experimental.explain import _fmt_row_count, explain_query +from cudf_polars.experimental.explain import ( + _fmt_row_count, + explain_query, + serialize_query, +) from cudf_polars.testing.asserts import ( DEFAULT_CLUSTER, DEFAULT_RUNTIME, @@ -17,6 +24,9 @@ ) from cudf_polars.testing.io import make_lazy_frame, make_partitioned_source +if TYPE_CHECKING: + from pathlib import Path + @pytest.fixture(scope="module") def df(): @@ -497,6 +507,153 @@ def _gb(df): ) +def test_serialize_query(): + # this test is sensitive to the polars version. + # we get a different query plan for polars < 1.35.0. + pytest.importorskip("polars", minversion="1.35.0") + + left = pl.LazyFrame({"a": ["a", "b", "a"], "b": [1, 2, 3]}) + right = pl.LazyFrame({"a": ["a", "b", "c"], "c": [4, 5, 6]}) + + q = ( + left.join(right, on="a", how="inner") + .group_by("a") + .agg(pl.col("b").sum(), pl.col("c").max()) + ) + engine = pl.GPUEngine(executor="streaming", raise_on_fail=True) + dag = serialize_query(q, engine) + + # We don't know the exact node IDs, but we can check the structure. + assert len(dag.roots) == 1 + node_types = sorted({x.type for x in dag.nodes.values()}) + assert node_types == ["DataFrameScan", "GroupBy", "Join", "Projection", "Select"] + assert len(dag.nodes) == 6 + assert len(dag.partition_info) == 6 + node_ids = set(dag.nodes) + + for node_id, node in dag.nodes.items(): + assert node.id == node_id + assert node_id in node_ids + assert set(node.children) <= node_ids + + match node.type: + case "DataFrameScan": + assert node.children == [] + assert node.schema == {"a": "STRING", "b": "INT64"} or node.schema == { + "a": "STRING", + "c": "INT64", + } + assert node_id not in dag.roots + + case "Projection": + assert len(node.children) == 1 + assert node.schema == {"b": "INT64", "c": "INT64", "a": "STRING"} + assert node.properties == {} + + case "GroupBy": + assert len(node.children) == 1 + assert node.schema == {"a": "STRING", "b": "INT64", "c": "INT64"} + assert node.properties == {"keys": ["a"]} + assert node_id not in dag.roots + + case "Select": + assert len(node.children) == 1 + assert node.schema == {"a": "STRING", "b": "INT64", "c": "INT64"} + assert node.properties == {"columns": ["a", "b", "c"]} + assert node_id in dag.roots + + case "Join": + assert len(node.children) == 2 + assert node.schema == {"a": "STRING", "b": "INT64", "c": "INT64"} + assert node.properties == { + "how": "Inner", + "left_on": ["a"], + "right_on": ["a"], + } + assert node_id not in dag.roots + + # smoke test to ensure that the output is JSON serializable + json.dumps(dataclasses.asdict(dag)) + + +def test_scan_properties(tmp_path: Path): + pl.DataFrame({"a": [1, 2, 3]}).write_parquet(tmp_path / "test.parquet") + + q = pl.scan_parquet(tmp_path / "test.parquet") + engine = pl.GPUEngine(executor="streaming", raise_on_fail=True) + dag = serialize_query(q, engine) + + # walk Union -> Scan + node = dag.nodes[dag.nodes[dag.roots[0]].children[0]] + assert node.type == "Scan" + assert node.properties == { + "paths": [str(tmp_path / "test.parquet")], + "typ": "parquet", + } + + +@pytest.mark.parametrize("descending", [False, True]) +def test_sort_properties(*, descending: bool): + q = pl.LazyFrame({"a": [1, 3, 2]}).sort("a", descending=descending) + dag = serialize_query(q, pl.GPUEngine(executor="streaming")) + + order = "DESCENDING" if descending else "ASCENDING" + node = dag.nodes[dag.roots[0]] + assert node.type == "Sort" + assert node.properties == {"by": ["a"], "order": [order]} + + +@pytest.mark.parametrize( + "predicate, expected", + [ + ( + pl.col("a") > 1, + { + "predicate": "a", + "op": "GREATER", + "left": {"type": "Col", "name": "a"}, + "right": {"type": "Literal", "value": 1}, + }, + ), + ( + pl.col("a") == pl.col("b"), + { + "predicate": "a", + "op": "EQUAL", + "left": {"type": "Col", "name": "a"}, + "right": {"type": "Col", "name": "b"}, + }, + ), + ], +) +def test_filter_properties(predicate: pl.Expr, expected: dict): + q = pl.LazyFrame({"a": [1, 2, 3], "b": [2, 2, 2]}).filter(predicate) + dag = serialize_query(q, pl.GPUEngine(executor="streaming")) + + node = dag.nodes[dag.roots[0]] + assert node.type == "Filter" + assert node.properties == expected + + +def test_select_properties(): + q = pl.LazyFrame({"a": [1, 2, 3]}).select(pl.col("a") + 1) + dag = serialize_query(q, pl.GPUEngine(executor="streaming")) + + node = dag.nodes[dag.roots[0]] + assert node.type == "Select" + assert node.properties == {"columns": ["a"]} + + +def test_hstack_properties(): + left = pl.LazyFrame({"a": [1, 2, 3]}) + q = left.with_columns(pl.col("a"), (pl.col("a") + 1).alias("b")) + dag = serialize_query(q, pl.GPUEngine(executor="streaming")) + + node = dag.nodes[dag.roots[0]] + assert node.type == "HStack" + assert node.properties == {"columns": ["a", "b"]} + + @pytest.mark.skipif( DEFAULT_RUNTIME != "rapidsmpf", reason="Requires the rapidsmpf runtime",