Improve scan partitioning planning to prefer sizes closer to target

[Matt711]

Description

For parquet scans, we use the target_partition_size(say T) and the total uncompressed size in bytes (say S) to decide whether to split a large file into smaller reads (ie. SplitScans) or fuse together smaller files into larger reads (ie. FusedScans). If S < T, then k, the number of splits is k = ceil(S / T). This means we always round up the number of splits. And so we run into the situation where we split a file into two parts when the entire file size was closer to T. In that situation, one task should read the entire file rather than splitting. Similarly for two and three splits, four and five splits, and so on. What we should do instead is decide whether to split into k or k+1 parts by how close the projected number of bytes is to T. And choose which ever projection is closer. If they tie, choose k.

For FusedScans, we should compare the projected number of bytes read for a fused task using k, the number of files fused, and k+1. And similarly choose whichever projection is closer to T. For ties, prefer k+1.

Overall, the idea is to get each "task" as close to reading a target_partition_size worth of bytes as possible.

Benchmarks

On 1xB200,

Benchmark  Branch            Avg lukewarm (s)  Avg hot (s)
TPC-DS     main              448.385           437.692
           this PR           395.437           386.210
TPC-H      main              131.884           128.187
           this PR           119.677           118.426

Around a 10% improvement for H/DS benchmarks at SF1K.

Follow up from #22945

Checklist

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

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[Matt711]

/ok to test 2f8bff8

[coderabbitai]

📝 Walkthrough

Walkthrough

scan_partition_plan in streaming/io.py replaces the previous split-or-fuse branching with a k_lo/k_hi nearest-blocksize heuristic that picks the factor whose chunk size is closest to the configured blocksize, returning SPLIT_FILES only when the chosen factor is ≥ 2. Tests add FooSource, FooStats, and _make_config helpers and a parameterized test validating the new factor and flavor selection.

Changes

Nearest-blocksize partition plan heuristic

Layer / File(s)	Summary
Nearest-blocksize factor selection python/cudf_polars/cudf_polars/streaming/io.py	Computes k_lo/k_hi candidate factors and selects whichever yields a chunk size closer to blocksize; returns SPLIT_FILES only when the chosen factor is ≥ 2, otherwise falls through to FUSED_FILES.
Test doubles and parameterized coverage python/cudf_polars/tests/streaming/test_scan.py	Adds FooSource, FooStats, _make_config, and test_scan_partition_plan_nearest to assert expected factor and flavor across multiple (file_size, n_paths) cases.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• rapidsai/cudf#22846: Modifies IOPartitionPlan.estimated_chunk_bytes plumbing that directly interacts with the factor selection changed here.
• rapidsai/cudf#22945: Also modifies scan_partition_plan factor/flavor selection in the same file.

Suggested reviewers

• madsbk

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 9.09% 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 summarizes the main change: scan partitioning now prefers task sizes closer to the target.
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.
Description check	✅ Passed	The description matches the code changes, explaining the new parquet scan partitioning heuristic and its test/benchmark context.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[madsbk]

Looks good

[madsbk]

Suggested change

	if file_size > blocksize:
	# Split large files into parts or fuse small files together, picking
	# whichever factor keeps each chunk closest to the target size.
	if file_size > blocksize:

[Matt711]

Will tack this on elsewhere

[vyasr]

Blocking until we can determine whether this improvement is effectively equivalent to just raising the target size.

[rjzamora]

Blocking until we can determine whether this improvement is effectively equivalent to just raising the target size.

Agree. To be concrete: Can you maybe collect results with the target_partition_size 20% larger?

[vyasr]

We had a call and basically concluded that the benchmarks we wanted are basically orthogonal to this PR. This PR is primarily changing how the target_partition_size parameter works, but we still have a tunable parameter, so the change just shifts the parameter space. If this is a better representation of what the parameter should mean, it makes sense to merge the PR regardless, then follow up with more measurements of how the partition size at different points in the query pipeline is potentially having a differential impact in case we should be tuning them differently.

[vyasr]

We had a call and basically concluded that the benchmarks we wanted are basically orthogonal to this PR. This PR is primarily changing how the target_partition_size parameter works, but we still have a tunable parameter, so the change just shifts the parameter space. If this is a better representation of what the parameter should mean, it makes sense to merge the PR regardless, then follow up with more measurements of how the partition size at different points in the query pipeline is potentially having a differential impact in case we should be tuning them differently.

[Matt711]

/merge