Add a I/O partition planning information to benchmark runner

[Matt711]

Description

Table: --explain-partition-plan

Partition Plan Summary
+----+----------+--------+--------------+-------+------------------------+-----------+------------+
| Q  | Table    | Flavor | Factor       | Files | Projected (bytes/file) | Size/task | Partitions |
+----+----------+--------+--------------+-------+------------------------+-----------+------------+
| 1  | lineitem | SPLIT  | 3 tasks/file | 60    | 4.35 GB                | 1.45 GB   | 180        |
| 5  | region   | FUSED  | 1 file/task  | 1     | 114 B                  | 114 B     | 1          |
|    | nation   | FUSED  | 1 file/task  | 1     | 329 B                  | 329 B     | 1          |
|    | customer | FUSED  | 1 file/task  | 2     | 900 MB                 | 900 MB    | 2          |
|    | orders   | FUSED  | 1 file/task  | 15    | 1.38 GB                | 1.38 GB   | 15         |
|    | lineitem | SPLIT  | 2 tasks/file | 60    | 2.71 GB                | 1.36 GB   | 120        |
|    | supplier | FUSED  | 1 file/task  | 1     | 120 MB                 | 120 MB    | 1          |
| 18 | orders   | SPLIT  | 2 tasks/file | 15    | 2.18 GB                | 1.09 GB   | 30         |
|    | lineitem | SPLIT  | 2 tasks/file | 60    | 1.6 GB                 | 800.1 MB  | 120        |
|    | customer | SPLIT  | 2 tasks/file | 2     | 1.95 GB                | 976.54 MB | 4          |
| 21 | lineitem | FUSED  | 1 file/task  | 60    | 1.11 GB                | 1.11 GB   | 60         |
|    | lineitem | FUSED  | 1 file/task  | 60    | 1.42 GB                | 1.42 GB   | 60         |
|    | supplier | FUSED  | 1 file/task  | 1     | 300.41 MB              | 300.41 MB | 1          |
|    | nation   | FUSED  | 1 file/task  | 1     | 259 B                  | 259 B     | 1          |
|    | orders   | FUSED  | 1 file/task  | 15    | 1.3 GB                 | 1.3 GB    | 15         |
+----+----------+--------+--------------+-------+------------------------+-----------+------------+

Explain output: CUDF_POLARS__EXPLAIN__PARTITION_PLAN=1 and --explain

SORT ('l_returnflag', 'l_linestatus') ('l_returnflag', 'l_linestatus', 'sum_qty', '...', 'avg_disc', 'count_order') [180]
  SELECT ('l_returnflag', 'l_linestatus', 'sum_qty', '...', 'avg_disc', 'count_order') [180]
    GROUPBY ('l_returnflag', 'l_linestatus') ('l_returnflag', 'l_linestatus', 'sum_qty', '...', 'avg_disc', '______________11') [180]
      HSTACK ('l_returnflag', 'l_linestatus', 'l_quantity', '...', 'l_tax', '__POLARS_CSER_0x6a4988e5b78df9ee') [180]
        PROJECTION ('l_returnflag', 'l_linestatus', 'l_quantity', 'l_extendedprice', 'l_discount', 'l_tax') [180]
          STREAMINGSCAN ('l_returnflag', 'l_linestatus', 'l_quantity', '...', 'l_tax', 'l_shipdate') [flavor=SPLIT_FILES factor=3 projected=4.35 GB] [180]

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds partition-plan collection and formatting for streaming queries, exposes the data through benchmark execution, appends IR annotations for streaming scans, and caps Parquet fusion factor by available paths.

Changes

Partition Plan Summary Feature

Layer / File(s)	Summary
PartitionPlanRow schema and IO factor fix python/cudf_polars/cudf_polars/streaming/explain.py, python/cudf_polars/cudf_polars/streaming/io.py	PartitionPlanRow defines the per-query partition-plan fields. explain.py adds path and environment helpers, and scan_partition_plan now caps Parquet fusion factor at len(ir.paths).
Partition plan collection and IR annotations python/cudf_polars/cudf_polars/streaming/explain.py	collect_partition_plan lowers the query, finds StreamingScan nodes, deduplicates scans, computes projected and task bytes, and infers table names. _repr_ir_tree accepts config and can append partition-plan metadata to StreamingScan headers when the explain flag is enabled.
Benchmark CLI integration python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py	Adds --explain-partition-plan, carries collected rows on QueryRunResult, accumulates rows across the query loop, and prints a combined partition-plan table after all queries finish.
Partition plan tests python/cudf_polars/tests/streaming/test_explain.py	Adds coverage for partition-plan collection, deduplication, table-name inference, table formatting, byte formatting, and factor wording.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Suggested reviewers

• mroeschke
• rjzamora

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 68.18% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly matches the main change by describing added I/O partition-planning info in the benchmark runner.
Description check	✅ Passed	The description directly covers the new partition-plan summary and explain-output integration.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

python/cudf_polars/tests/streaming/test_explain.py (1)

703-743: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoff

Consider adding edge case tests for partition plan collection.

The current tests cover the main FUSED_FILES and SPLIT_FILES scenarios well, but miss important edge cases:

• Empty DataFrame input
• All-null columns (affects projected_bytes calculation)
• Zero files case

These edge cases could reveal bugs in byte calculations, partition logic, or division-by-zero scenarios in the factor arithmetic.

Consider adding tests like:

def test_collect_partition_plan_empty_dataframe(tmp_path):
    """Empty DataFrame should handle gracefully."""
    empty_df = pl.DataFrame({"x": []})
    empty_df.write_parquet(tmp_path / "empty.parquet")
    q = pl.scan_parquet(tmp_path / "empty.parquet")
    engine = pl.GPUEngine(executor="streaming", raise_on_fail=True)
    rows = collect_partition_plan(q, engine, q_id=1)
    # Assert expected behavior for empty input

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@python/cudf_polars/tests/streaming/test_explain.py` around lines 703 - 743,
Add three new test functions to cover edge cases in partition plan collection:
test_collect_partition_plan_empty_dataframe to handle empty DataFrames,
test_collect_partition_plan_all_null_columns to test all-null column scenarios
that may affect projected_bytes calculations, and
test_collect_partition_plan_zero_files for the zero files case. Each test should
follow the same structure as the existing test_collect_partition_plan_fused and
test_collect_partition_plan_split functions, using make_partitioned_source to
set up the data and collect_partition_plan to verify the expected behavior
without errors or division-by-zero issues.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@python/cudf_polars/tests/streaming/test_explain.py`:
- Around line 703-743: Add three new test functions to cover edge cases in
partition plan collection: test_collect_partition_plan_empty_dataframe to handle
empty DataFrames, test_collect_partition_plan_all_null_columns to test all-null
column scenarios that may affect projected_bytes calculations, and
test_collect_partition_plan_zero_files for the zero files case. Each test should
follow the same structure as the existing test_collect_partition_plan_fused and
test_collect_partition_plan_split functions, using make_partitioned_source to
set up the data and collect_partition_plan to verify the expected behavior
without errors or division-by-zero issues.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 4e77b53a-8220-48cc-9cc3-47dc9986b967

📥 Commits

Reviewing files that changed from the base of the PR and between 9d8b246 and 6a8e16e.

📒 Files selected for processing (4)

• python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py
• python/cudf_polars/cudf_polars/streaming/explain.py
• python/cudf_polars/cudf_polars/streaming/io.py
• python/cudf_polars/tests/streaming/test_explain.py

[vyasr]

Was there a particular request that this PR was created for, or just something you thought would be useful?

[Matt711]

/merge

[Matt711]

Was there a particular request that this PR was created for, or just something you thought would be useful?

Yes I'm currently studying how well our split / fused task scans are aligned with the target partition size. I think we're currently being too conservative. I'll have a PR up soon.