Add a I/O partition planning information to benchmark runner

python/cudf_polars/cudf_polars/streaming/benchmarks/utils.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
 
 """Benchmark utilities for the RapidsMPF SPMD and Ray frontends."""
@@ -263,6 +263,7 @@ class QueryRunResult:
     plan: SerializablePlan | None
     iteration_failures: list[tuple[int, int]]
     validation_failed: bool
+    partition_plan_rows: list = dataclasses.field(default_factory=list)
 
 
 @dataclasses.dataclass
@@ -991,6 +992,16 @@ def run_polars_query(
 
         plan = serialize_query(q, engine)
 
+    part_plan_rows = []
+    if (
+        getattr(args, "explain_partition_plan", False)
+        and engine is not None
+        and run_config.frontend in _STREAMING_FRONTENDS
+    ):
+        from cudf_polars.streaming.explain import collect_partition_plan
+
+        part_plan_rows = collect_partition_plan(q, engine, q_id)
+
     casts = benchmark.EXPECTED_CASTS.get(q_id, [])
     if numeric_type == "decimal":
         casts.extend(benchmark.EXPECTED_CASTS_DECIMAL.get(q_id, []))
@@ -1086,6 +1097,7 @@ def run_polars_query(
         plan=plan,
         iteration_failures=iteration_failures,
         validation_failed=validation_failed,
+        partition_plan_rows=part_plan_rows,
     )
 
 
@@ -1108,6 +1120,7 @@ def _run_query_loop(
     plans: dict[int, SerializablePlan] = {}
     validation_failures: list[int] = []
     query_failures: list[tuple[int, int]] = []
+    all_partition_plan_rows: list = []
 
     for q_id in run_config.queries:
         try:
@@ -1143,6 +1156,12 @@ def _run_query_loop(
         query_failures.extend(result.iteration_failures)
         if result.validation_failed:
             validation_failures.append(q_id)
+        all_partition_plan_rows.extend(result.partition_plan_rows)
+
+    if all_partition_plan_rows and getattr(args, "explain_partition_plan", False):
+        from cudf_polars.streaming.explain import format_partition_plan_table
+
+        print(format_partition_plan_table(all_partition_plan_rows), flush=True)
 
     return records, plans, validation_failures, query_failures
 
@@ -1945,6 +1964,12 @@ def build_parser(num_queries: int = 22) -> argparse.ArgumentParser:
         help="Print an outline of the logical plan.",
         default=False,
     )
+    parser.add_argument(
+        "--explain-partition-plan",
+        action=argparse.BooleanOptionalAction,
+        help="Print a combined partition plan summary table across all queries.",
+        default=False,
+    )
     parser.add_argument(
         "--print-plans",
         action=argparse.BooleanOptionalAction,

python/cudf_polars/cudf_polars/streaming/explain.py

@@ -9,9 +9,11 @@
 import dataclasses
 import datetime
 import functools
+import os
 import os.path
 from collections.abc import Mapping, Sequence
 from itertools import groupby
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, Self, TypeAlias
 
 import pylibcudf as plc
@@ -33,7 +35,8 @@
 )
 from cudf_polars.dsl.translate import Translator
 from cudf_polars.dsl.traversal import traversal
-from cudf_polars.streaming.io import StreamingScan
+from cudf_polars.streaming.base import IOPartitionFlavor
+from cudf_polars.streaming.io import StreamingScan, scan_partition_plan
 from cudf_polars.streaming.parallel import lower_ir_graph
 from cudf_polars.streaming.shuffle import Shuffle
 from cudf_polars.streaming.statistics import (
@@ -51,6 +54,20 @@
     from cudf_polars.streaming.base import PartitionInfo, StatsCollector
 
 
+@dataclasses.dataclass
+class PartitionPlanRow:
+    """One row of the partition plan summary table."""
+
+    query: int
+    table: str
+    flavor: IOPartitionFlavor
+    factor: int
+    files: int
+    projected_bytes: int
+    task_bytes: int
+    partitions: int
+
+
 Serializable: TypeAlias = (
     str
     | int
@@ -106,7 +123,7 @@ def explain_query(
         with cm:
             stats = collect_statistics(ir, config, executor)
         lowered_ir, partition_info = lower_ir_graph(ir, config, stats)
-        return _repr_ir_tree(lowered_ir, partition_info)
+        return _repr_ir_tree(lowered_ir, partition_info, stats=stats, config=config)
     else:
         if config.executor.name == "streaming":
             # Include row-count statistics for the logical plan
@@ -117,6 +134,162 @@ def explain_query(
             return _repr_ir_tree(ir)
 
 
+def collect_partition_plan(
+    q: pl.LazyFrame,
+    engine: pl.GPUEngine,
+    q_id: int,
+) -> list[PartitionPlanRow]:
+    """
+    Return one PartitionPlanRow per unique StreamingScan in the physical plan.
+
+    Deduplicates scans that appear multiple times due to subquery structure.
+    """
+    config = ConfigOptions.from_polars_engine(engine)
+    ir = Translator(q._ldf.visit(), engine).translate_ir()
+
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        stats = collect_statistics(ir, config, executor)
+    lowered_ir, partition_info = lower_ir_graph(ir, config, stats)
+
+    seen: set[tuple] = set()
+    rows: list[PartitionPlanRow] = []
+
+    for node in traversal([lowered_ir]):
+        if not isinstance(node, StreamingScan):
+            continue
+        base_scan = node.base_scan
+
+        dedup_key = (tuple(base_scan.paths), tuple(sorted(base_scan.schema.keys())))
+        if dedup_key in seen:
+            continue
+        seen.add(dedup_key)
+
+        source = stats.scan_stats.get(base_scan)
+        if source is None:
+            continue
+
+        plan = scan_partition_plan(base_scan, stats, config)
+        projected_bytes = sum(
+            sz
+            for col in base_scan.schema
+            if (sz := source.column_storage_size(col)) is not None
+        )
+        partitions = partition_info[node].count
+        factor = plan.factor
+        flavor = plan.flavor
+
+        match flavor:
+            case IOPartitionFlavor.SPLIT_FILES:
+                files = partitions // factor if factor > 0 else partitions
+                task_bytes = (
+                    projected_bytes // factor if factor > 0 else projected_bytes
+                )
+            case IOPartitionFlavor.FUSED_FILES:
+                files = partitions * factor
+                task_bytes = projected_bytes * factor
+            case _:
+                files = partitions
+                task_bytes = projected_bytes
+
+        p = Path(base_scan.paths[0])
+        stem = p.stem
+        parent = p.parent.name
+        # Prefer the stem unless it looks like a partition filename (purely
+        # numeric like "1" or prefixed like "part-0"), in which case the
+        # parent directory holds the table name.
+        table = parent if (stem.isdigit() or stem.lower().startswith("part")) else stem
+
+        rows.append(
+            PartitionPlanRow(
+                query=q_id,
+                table=table,
+                flavor=flavor,
+                factor=factor,
+                files=files,
+                projected_bytes=projected_bytes,
+                task_bytes=task_bytes,
+                partitions=partitions,
+            )
+        )
+
+    return rows
+
+
+def _fmt_partition_bytes(b: int) -> str:
+    if b < 1_000:
+        return f"{b} B"
+    elif b < 1_000_000:
+        return f"{round(b / 1_000, 2):g} KB"
+    elif b < 1_000_000_000:
+        return f"{round(b / 1_000_000, 2):g} MB"
+    else:
+        return f"{round(b / 1_000_000_000, 2):g} GB"
+
+
+def factor_str(row: PartitionPlanRow) -> str:
+    """Format the factor field with units appropriate to the scan flavor."""
+    match row.flavor:
+        case IOPartitionFlavor.SPLIT_FILES:
+            return f"{row.factor} tasks/file"
+        case IOPartitionFlavor.FUSED_FILES:
+            unit = "file" if row.factor == 1 else "files"
+            return f"{row.factor} {unit}/task"
+        case _:
+            return str(row.factor)
+
+
+def format_partition_plan_table(rows: list[PartitionPlanRow]) -> str:
+    """Format a list of PartitionPlanRows as a fixed-width ASCII table."""
+    if not rows:
+        return ""
+
+    headers = [
+        "Q",
+        "Table",
+        "Flavor",
+        "Factor",
+        "Files",
+        "Projected (bytes/file)",
+        "Size/task",
+        "Partitions",
+    ]
+
+    formatted: list[list[str]] = []
+    prev_q: int | None = None
+    for row in rows:
+        q_str = str(row.query) if row.query != prev_q else ""
+        prev_q = row.query
+        formatted.append(
+            [
+                q_str,
+                row.table,
+                row.flavor.name,
+                factor_str(row),
+                str(row.files),
+                _fmt_partition_bytes(row.projected_bytes),
+                _fmt_partition_bytes(row.task_bytes),
+                str(row.partitions),
+            ]
+        )
+
+    col_widths = [len(h) for h in headers]
+    for cells in formatted:
+        for i, cell in enumerate(cells):
+            col_widths[i] = max(col_widths[i], len(cell))
+
+    sep = "+-" + "-+-".join("-" * w for w in col_widths) + "-+"
+    header_row = (
+        "| " + " | ".join(h.ljust(col_widths[i]) for i, h in enumerate(headers)) + " |"
+    )
+    lines = ["", "Partition Plan Summary", sep, header_row, sep]
+    lines.extend(
+        "| " + " | ".join(c.ljust(col_widths[i]) for i, c in enumerate(cells)) + " |"
+        for cells in formatted
+    )
+    lines.append(sep)
+    return "\n".join(lines)
+
+
 def serialize_query(
     q: pl.LazyFrame,
     engine: pl.GPUEngine,
@@ -207,6 +380,7 @@ def _repr_ir_tree(
     *,
     offset: str = "",
     stats: StatsCollector | None = None,
+    config: ConfigOptions | None = None,
 ) -> str:
     header = _repr_ir(ir, offset=offset)
     count = partition_info[ir].count if partition_info else None
@@ -215,11 +389,31 @@ def _repr_ir_tree(
         row_count_estimate = _fmt_row_count(source.row_count)
         row_count = f"~{row_count_estimate}" if row_count_estimate else "unknown"
         header = header.rstrip("\n") + f" {row_count=}\n"
+    if (
+        os.environ.get("CUDF_POLARS__EXPLAIN__PARTITION_PLAN", "0") == "1"
+        and config is not None
+        and stats is not None
+        and isinstance(ir, StreamingScan)
+        and (source := stats.scan_stats.get(ir.base_scan)) is not None
+    ):
+        plan = scan_partition_plan(ir.base_scan, stats, config)
+        projected_size = sum(
+            sz
+            for col in ir.base_scan.schema
+            if (sz := source.column_storage_size(col)) is not None
+        )
+        plan_info = (
+            f"flavor={plan.flavor.name} factor={plan.factor}"
+            f" projected={_fmt_partition_bytes(projected_size)}"
+        )
+        header = header.rstrip("\n") + f" [{plan_info}]\n"
     if count is not None:
         header = header.rstrip("\n") + f" [{count}]\n"
 
     children_strs = [
-        _repr_ir_tree(child, partition_info, offset=offset + "  ", stats=stats)
+        _repr_ir_tree(
+            child, partition_info, offset=offset + "  ", stats=stats, config=config
+        )
         for child in ir.children
     ]
 

python/cudf_polars/cudf_polars/streaming/io.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 """Multi-partition IO Logic."""
 
@@ -98,7 +98,7 @@ def scan_partition_plan(
                     )
                 else:
                     # Fuse small files
-                    factor = max(blocksize // int(file_size), 1)
+                    factor = min(max(blocksize // int(file_size), 1), len(ir.paths))
                     return IOPartitionPlan(
                         factor,
                         IOPartitionFlavor.FUSED_FILES,