Add coarse-grained streaming-node tracing

docs/cudf/source/cudf_polars/api.md

@@ -11,6 +11,7 @@ For the most part, the public API of `cudf-polars` is the polars API.
       DynamicPlanningOptions,
       InMemoryExecutor,
       ParquetOptions,
+      TracingOptions,
       Cluster,
       ShuffleMethod,
       ShufflerInsertionMethod,

python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py

@@ -257,6 +257,7 @@ class RunConfig:
     collect_traces: bool = False
     stats_planning: bool
     dynamic_planning: bool | None = None
+    trace_output_path: str | None = None
     max_io_threads: int
     native_parquet: bool
     spill_to_pinned_memory: bool
@@ -377,6 +378,7 @@ def from_args(cls, args: argparse.Namespace) -> RunConfig:
             collect_traces=args.collect_traces,
             stats_planning=args.stats_planning,
             dynamic_planning=args.dynamic_planning,
+            trace_output_path=args.trace_output_path,
             max_io_threads=args.max_io_threads,
             native_parquet=args.native_parquet,
             extra_info=args.extra_info,
@@ -477,6 +479,8 @@ def get_executor_options(
         if run_config.dynamic_planning:
             # Pass empty dict to enable with defaults; None means disabled
             executor_options["dynamic_planning"] = {}
+        if run_config.trace_output_path is not None:
+            executor_options["tracing"] = {"output_path": run_config.trace_output_path}
 
     if (
         benchmark
@@ -974,6 +978,13 @@ def parse_args(
         default=False,
         help="Enable dynamic shuffle planning (not yet implemented). ",
     )
+    parser.add_argument(
+        "--trace-output-path",
+        dest="trace_output_path",
+        type=str,
+        default=None,
+        help="Path to write tracing output (row counts per node).",
+    )
     parser.add_argument(
         "--max-io-threads",
         default=2,

python/cudf_polars/cudf_polars/experimental/explain.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
@@ -26,11 +26,13 @@
 
 if TYPE_CHECKING:
     from collections.abc import MutableMapping
+    from pathlib import Path
 
     import polars as pl
 
     from cudf_polars.dsl.ir import IR
     from cudf_polars.experimental.base import PartitionInfo, StatsCollector
+    from cudf_polars.experimental.rapidsmpf.tracing import StreamingQueryTracer
 
 
 def explain_query(
@@ -168,3 +170,67 @@ 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)
+
+
+def write_query_trace(
+    trace_output: str | Path,
+    ir: IR,
+    partition_info: MutableMapping[IR, PartitionInfo],
+    tracer: StreamingQueryTracer,
+) -> None:
+    """
+    Write a post-execution trace showing actual row counts and decisions.
+
+    Parameters
+    ----------
+    trace_output
+        Path to write the trace file.
+    ir
+        The lowered IR root node.
+    partition_info
+        Partition information for the IR nodes.
+    tracer
+        The tracer with actual row counts and decisions from execution.
+    """
+    from pathlib import Path
+
+    trace_repr = _repr_trace_tree(ir, partition_info, tracer)
+    Path(trace_output).write_text(trace_repr)
+
+
+def _repr_trace_tree(
+    ir: IR,
+    partition_info: MutableMapping[IR, PartitionInfo],
+    tracer: StreamingQueryTracer,
+    *,
+    offset: str = "",
+) -> str:
+    """Recursively build a tree representation with tracer data."""
+    header = _repr_ir(ir, offset=offset)
+    header = header.rstrip("\n")
+
+    # Get node tracer if it exists
+    if (node_tracer := tracer.node_tracers.get(ir)) is not None:
+        # Add actual row count if available
+        if node_tracer.row_count is not None:
+            header += f" rows={_fmt_row_count(node_tracer.row_count)}"
+
+        # Add decision if present
+        if node_tracer.decision is not None:
+            header += f" decision={node_tracer.decision}"
+
+        # Add actual chunk count
+        header += f" chunks={node_tracer.chunk_count}"
+
+    children_strs = [
+        _repr_trace_tree(child, partition_info, tracer, offset=offset + "  ")
+        for child in ir.children
+    ]
+
+    header += "\n"
+    return header + "".join(
+        f"{line}{offset}  (repeated {count} times)\n"
+        if (count := sum(1 for _ in group)) > 1
+        else line
+        for line, group in groupby(children_strs)
+    )

python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py

@@ -32,13 +32,15 @@
 import cudf_polars.experimental.rapidsmpf.union  # noqa: F401
 from cudf_polars.containers import DataFrame
 from cudf_polars.dsl.ir import DataFrameScan, IRExecutionContext, Join, Scan, Union
+from cudf_polars.dsl.tracing import LOG_TRACES
 from cudf_polars.dsl.traversal import CachingVisitor, traversal
 from cudf_polars.experimental.rapidsmpf.collectives import ReserveOpIDs
 from cudf_polars.experimental.rapidsmpf.dispatch import FanoutInfo, lower_ir_node
 from cudf_polars.experimental.rapidsmpf.nodes import (
     generate_ir_sub_network_wrapper,
     metadata_drain_node,
 )
+from cudf_polars.experimental.rapidsmpf.tracing import StreamingQueryTracer
 from cudf_polars.experimental.rapidsmpf.utils import empty_table_chunk
 from cudf_polars.experimental.statistics import collect_statistics
 from cudf_polars.experimental.utils import _concat
@@ -107,7 +109,7 @@ def evaluate_logical_plan(
             # NOTE: Distributed execution requires Dask for now
             from cudf_polars.experimental.rapidsmpf.dask import evaluate_pipeline_dask
 
-            result, metadata_collector = evaluate_pipeline_dask(
+            result, metadata_collector, tracer = evaluate_pipeline_dask(
                 evaluate_pipeline,
                 ir,
                 partition_info,
@@ -118,7 +120,7 @@ def evaluate_logical_plan(
             )
         else:
             # Single-process execution: Run locally
-            result, metadata_collector = evaluate_pipeline(
+            result, metadata_collector, tracer = evaluate_pipeline(
                 ir,
                 partition_info,
                 config_options,
@@ -127,6 +129,13 @@ def evaluate_logical_plan(
                 collect_metadata=collect_metadata,
             )
 
+    # Write tracer output if configured
+    tracing = config_options.executor.tracing
+    if tracing is not None and tracing.output_path is not None and tracer is not None:
+        from cudf_polars.experimental.explain import write_query_trace
+
+        write_query_trace(tracing.output_path, ir, partition_info, tracer)
+
     return result, metadata_collector
 
 
@@ -139,7 +148,7 @@ def evaluate_pipeline(
     rmpf_context: Context | None = None,
     *,
     collect_metadata: bool = False,
-) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]:
+) -> tuple[pl.DataFrame, list[ChannelMetadata] | None, StreamingQueryTracer | None]:
     """
     Build and evaluate a RapidsMPF streaming pipeline.
 
@@ -162,7 +171,7 @@ def evaluate_pipeline(
 
     Returns
     -------
-    The output DataFrame and metadata collector.
+    The output DataFrame, metadata collector, and tracer.
     """
     assert config_options.executor.name == "streaming", "Executor must be streaming"
     assert config_options.executor.runtime == "rapidsmpf", "Runtime must be rapidsmpf"
@@ -232,6 +241,11 @@ def evaluate_pipeline(
         metadata_collector: list[ChannelMetadata] | None = (
             [] if collect_metadata else None
         )
+        tracer: StreamingQueryTracer | None = (
+            StreamingQueryTracer()
+            if config_options.executor.tracing is not None or LOG_TRACES
+            else None
+        )
         nodes, output = generate_network(
             rmpf_context,
             ir,
@@ -241,6 +255,7 @@ def evaluate_pipeline(
             ir_context=ir_context,
             collective_id_map=collective_id_map,
             metadata_collector=metadata_collector,
+            tracer=tracer,
         )
 
         # Run the network
@@ -294,7 +309,7 @@ def evaluate_pipeline(
         if _initial_mr is not None:
             rmm.mr.set_current_device_resource(_original_mr)
 
-        return result, metadata_collector
+        return result, metadata_collector, tracer
 
 
 def lower_ir_graph(
@@ -416,6 +431,7 @@ def generate_network(
     ir_context: IRExecutionContext,
     collective_id_map: dict[IR, list[int]],
     metadata_collector: list[ChannelMetadata] | None,
+    tracer: StreamingQueryTracer | None = None,
 ) -> tuple[list[Any], DeferredMessages]:
     """
     Translate the IR graph to a RapidsMPF streaming network.
@@ -440,6 +456,8 @@ def generate_network(
         The list to collect the final metadata.
         This list will be mutated when the network is executed.
         If None, metadata will not be collected.
+    tracer
+        Profiler for collecting runtime statistics.
 
     Returns
     -------
@@ -472,6 +490,7 @@ def generate_network(
         "max_io_threads": max_io_threads_local,
         "stats": stats,
         "collective_id_map": collective_id_map,
+        "tracer": tracer,
     }
     mapper: SubNetGenerator = CachingVisitor(
         generate_ir_sub_network_wrapper, state=state

python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py

@@ -24,6 +24,7 @@
     from cudf_polars.dsl.ir import IR
     from cudf_polars.experimental.base import PartitionInfo, StatsCollector
     from cudf_polars.experimental.parallel import ConfigOptions
+    from cudf_polars.experimental.rapidsmpf.tracing import StreamingQueryTracer
 
 
 class EvaluatePipelineCallback(Protocol):
@@ -39,8 +40,8 @@ def __call__(
         rmpf_context: Context | None = None,
         *,
         collect_metadata: bool = False,
-    ) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]:
-        """Evaluate a pipeline and return the result DataFrame and metadata."""
+    ) -> tuple[pl.DataFrame, list[ChannelMetadata] | None, StreamingQueryTracer | None]:
+        """Evaluate a pipeline and return the result DataFrame, metadata, and tracer."""
         ...
 
 
@@ -61,7 +62,7 @@ def evaluate_pipeline_dask(
     collective_id_map: dict[IR, list[int]],
     *,
     collect_metadata: bool = False,
-) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]:
+) -> tuple[pl.DataFrame, list[ChannelMetadata] | None, StreamingQueryTracer | None]:
     """
     Evaluate a RapidsMPF streaming pipeline on a Dask cluster.
 
@@ -84,7 +85,7 @@ def evaluate_pipeline_dask(
 
     Returns
     -------
-    The output DataFrame and metadata collector.
+    The output DataFrame, metadata collector, and merged tracer.
     """
     client = get_dask_client()
     result = client.run(
@@ -99,12 +100,18 @@ def evaluate_pipeline_dask(
     )
     dfs: list[pl.DataFrame] = []
     metadata_collector: list[ChannelMetadata] = []
-    for df, md in result.values():
+    merged_tracer: StreamingQueryTracer | None = None
+    for df, md, tracer in result.values():
         dfs.append(df)
         if md is not None:
             metadata_collector.extend(md)
+        if tracer is not None:
+            if merged_tracer is None:
+                merged_tracer = tracer
+            else:
+                merged_tracer.merge(tracer)
 
-    return pl.concat(dfs), metadata_collector or None
+    return pl.concat(dfs), metadata_collector or None, merged_tracer
 
 
 def _evaluate_pipeline_dask(
@@ -117,7 +124,7 @@ def _evaluate_pipeline_dask(
     dask_worker: Any = None,
     *,
     collect_metadata: bool = False,
-) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]:
+) -> tuple[pl.DataFrame, list[ChannelMetadata] | None, StreamingQueryTracer | None]:
     """
     Build and evaluate a RapidsMPF streaming pipeline.
 
@@ -144,7 +151,7 @@ def _evaluate_pipeline_dask(
 
     Returns
     -------
-    The output DataFrame and metadata collector.
+    The output DataFrame, metadata collector, and tracer.
     """
     assert dask_worker is not None, "Dask worker must be provided"
     assert config_options.executor.name == "streaming", "Executor must be streaming"

python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py

@@ -15,10 +15,8 @@
     from rapidsmpf.streaming.core.context import Context
 
     from cudf_polars.dsl.ir import IR, IRExecutionContext
-    from cudf_polars.experimental.base import (
-        PartitionInfo,
-        StatsCollector,
-    )
+    from cudf_polars.experimental.base import PartitionInfo, StatsCollector
+    from cudf_polars.experimental.rapidsmpf.tracing import StreamingQueryTracer
     from cudf_polars.experimental.rapidsmpf.utils import ChannelManager
     from cudf_polars.utils.config import ConfigOptions
 
@@ -77,6 +75,8 @@ class GenState(TypedDict):
         Statistics collector.
     collective_id_map
         The mapping of IR nodes to lists of collective IDs.
+    tracer
+        Runtime tracer for collecting execution statistics.
     """
 
     context: Context
@@ -87,6 +87,7 @@ class GenState(TypedDict):
     max_io_threads: int
     stats: StatsCollector
     collective_id_map: dict[IR, list[int]]
+    tracer: StreamingQueryTracer | None
 
 
 SubNetGenerator: TypeAlias = GenericTransformer[