Capture query plan in pds-h output

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.

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: