Use conservative column-size estimates in cudf-polars

[rjzamora]

Description

While experimenting with some real user data, I discovered that our existing logic for detecting "suspicious" column-size estimates for Parquet data is sometimes insufficient. More specifically, I was running into OOM errors as a result of Scan chunks that were multiple times larger than target_partition_size would suggest.

This PR sets a more-conservative floor on the column-size estimates we use for Scan partition planning.

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

This PR adds a dtype-aware decoded-size floor helper, updates ParquetSourceInfo sizing to use per-column dtype floors and row-group samples, threads Parquet schema through cached source-info construction, accumulates schema in statistics grouping, and adds a unit test validating the floor application.

Changes

Dtype-aware Parquet sizing and schema threading

Layer / File(s)	Summary
Dtype sizing helper and source info update python/cudf_polars/cudf_polars/streaming/io.py	Adds _decoded_size_floor() mapping plc.TypeId/DataType to conservative per-row byte floors (including nullmask and variable-width types). ParquetSourceInfo.from_paths() signature now accepts schema: tuple[tuple[str, DataType], ...] and computes per_file_means[col] using max(footer_mean, decoded_floor) or max(footer_mean, decoded_floor, sampled_rowgroup_estimate) when samples exist.
Cache key and source-info plumbing python/cudf_polars/cudf_polars/streaming/io.py	_build_parquet_source() cache key and call updated to include schema; _build_source_info() accepts optional schema and forwards IR-derived tuple(ir.schema.items()) when not provided.
Statistics collection schema accumulation python/cudf_polars/cudf_polars/streaming/statistics.py	collect_statistics groups parquet scans by paths and accumulates a merged schema per group; job submissions pass tuple(schema.items()) along with needed_cols to _build_source_info().
Test for decoded dtype floor python/cudf_polars/tests/streaming/test_stats.py	Adds test_parquet_source_info_uses_decoded_dtype_floor which monkeypatches ParquetMetadata and _sample_rg_sizes, invokes ParquetSourceInfo.from_paths() with varied dtypes, and asserts expected column_storage_size results.

---

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

---

Suggested reviewers

• madsbk
• Matt711

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 41.67% 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 accurately describes the main change: implementing conservative column-size estimate floors for Parquet scan planning in cudf-polars.
Description check	✅ Passed	The description clearly explains the motivation (OOM errors from undersized estimates), the solution (conservative floor), and relates directly to the changeset.
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 and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 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.

Inline comments:
In `@python/cudf_polars/cudf_polars/streaming/io.py`:
- Around line 820-847: The _decoded_size_floor function omits several
fixed-width TypeIds so TIMESTAMP_SECONDS, DURATION_DAYS, DURATION_SECONDS,
DECIMAL32, and DECIMAL64 fall through to max(1, nrows) and under-estimate sizes;
update _decoded_size_floor to add explicit branches for these plc.TypeId values
and return nrows * <bytes> + nullmask using the same byte widths convention as
the surrounding cases (TIMESTAMP_SECONDS -> 8, DURATION_DAYS -> 8,
DURATION_SECONDS -> 8, DECIMAL32 -> 4, DECIMAL64 -> 8) so ParquetSourceInfo
per_file_means/file_size calculations used by scan_partition_plan are correct.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: de2531a5-7463-42f4-944f-9fac6e54d133

📥 Commits

Reviewing files that changed from the base of the PR and between 5d818ad and 5d76127.

📒 Files selected for processing (3)

• python/cudf_polars/cudf_polars/streaming/io.py
• python/cudf_polars/cudf_polars/streaming/statistics.py
• python/cudf_polars/tests/streaming/test_stats.py

[Matt711]

LGTM. This doesn't seem like it would negatively impact out nightly TPC-H/DS benchmarks. Probably worth keeping an eye on them though once this merges

[rjzamora]

This doesn't seem like it would negatively impact out nightly TPC-H/DS benchmarks. Probably worth keeping an eye on them though once this merges

This is actually an important consideration. I just started collecting TPC-H numbers on H100s, and this PR does affect benchmark performance in a way that is probably not surprising: Since some queries now use more partitions (because we were previously under-estimating the in-memory size of some columns), performance with the default target_partition_size is indeed slightly worse. For this benchmark, the underestimation on main isn't enough to cause a stability/OOM problem.

[rjzamora]

/merge