test(parquet): LayoutTest coverage for dictionary + FixedSizeBinary page sizing - #16
Merged
adriangb merged 7 commits intoJun 3, 2026
Conversation
Add the regression tests first, before any fix (TDD). With unmodified `main` the page-bounding assertions fail: the column writer only checks the data/dictionary page byte limit *after* each `write_batch_size` mini-batch, so large variable-width values pile into a single oversized page (we've observed 2 GiB data pages and ~256x dictionary-page overshoot at default settings). Column-writer tests (`ColumnValueEncoderImpl` path): - large BYTE_ARRAY values cap data pages near one value each - large values inside a repeated/list column (record-boundary stepping) - nullable column (value vs level counting) - dictionary spill then plain-encode large values - large distinct values bound the dictionary page - FIXED_LEN_BYTE_ARRAY byte budget Arrow-writer tests (`ByteArrayEncoder` path, what real users hit): - large `Utf8` strings via `ArrowWriter` - mixed small/large strings round-trip bit-identically - large `Utf8View` strings - all-null string column stays correct The subsequent commits make each of these pass. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Make the mini-batch size byte-budget aware so the post-write page check
fires before a data page grows unbounded:
- New `ColumnValueEncoder::count_values_within_byte_budget{,_gather}`
trait methods (default `None` = "no estimate; stay batched"), with a
concrete impl on `ColumnValueEncoderImpl` driven by
`plain_encoded_byte_size`. Fixed-width physical types answer in one
division; only variable-width BYTE_ARRAY/FLBA walk values, exiting at
the first that overruns the budget.
- New `LevelDataRef::value_count` converts a chunk's level span into a
leaf-value count (O(1) for flat columns, def-level scan for
nullable/nested), with a unit test.
- New `ByteBudgetChunker` picks the largest sub-batch that fits one page
budget. For the common case (small or fixed-width values) it returns
the full chunk with no value inspection, so the hot path is unchanged.
- `write_batch_internal` consults the chunker per chunk and, only when a
chunk would overflow, routes through the new `write_granular_chunk`,
which sub-batches so the post-write check fires in time.
Repeated/nested columns step on record (rep == 0) boundaries so a
record never spans pages.
This makes the column-writer data-page, list, nullable and FLBA
regression tests pass. Dictionary-encoding columns are still left on the
batched path (the data page holds only small RLE indices) — bounding the
dictionary page is a separate commit, so the two dictionary tests and
the arrow `ByteArrayEncoder` tests do not pass yet.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Implement `ColumnValueEncoder::count_values_within_byte_budget_gather` for `ByteArrayEncoder`, the encoder real `ArrowWriter` users hit. This makes the page-size bound fire for arrow string/binary columns; the previous commit only wired the generic column-writer path. Makes the arrow-writer regression tests pass. The implementation stays off the hot path for small values via cheap O(1) upper bounds before any per-value walk: - Offset-backed arrays (`Utf8`/`LargeUtf8`/`Binary`/`LargeBinary`): the span `offsets[last+1] - offsets[first]` bounds the chunk's payload in O(1); if it fits, every value fits. The span is exact even for nullable columns (skipped positions are nulls with zero offset delta), so sparse `indices` skip the per-value walk too. - View arrays (`Utf8View`/`BinaryView`): lengths live in the low 32 bits of each view word, so an O(1) `n * (max_value_len + 4)` bound skips the scan in the common case; otherwise scan lengths with no data-buffer deref. - Dictionary input: treated as always-fitting — dict-encoded arrow input implies values small enough to dedup, the opposite of the blob case this targets, and a per-key walk measurably regressed the bench. - FixedSizeBinary: falls through to the generic accessor walk. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
While a column dictionary-encodes, the data page holds only small RLE indices but the *dictionary* page accumulates the distinct values themselves, and its spill check runs only once per mini-batch. A mini-batch of large distinct values therefore interns `write_batch_size * value_size` bytes into the dictionary page before the check fires — ~256x the limit in the worst case. Extend `ByteBudgetChunker` to bound the dictionary-encoding phase too: when `has_dictionary()`, size the mini-batch against the dictionary page's *remaining* budget (`dict_page_byte_limit - estimated_dict_page_size`) rather than the data page. Fixed-width columns short-circuit via a precomputed `static_dict_always_fits`, so only large variable-width distinct values pay the per-value walk. Makes the two dictionary regression tests pass. `arrow_writer_layout`'s `test_string` updates accordingly: the dictionary page is now bounded at its 1000-byte limit and spills one mini-batch earlier (125 rows rather than 130) instead of overshooting to 1040. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The variable-width byte-budget walks returned the largest count whose cumulative encoded size was *under* the budget, so each mini-batch ended just short of the page threshold. When the input row batch did not divide evenly into mini-batches, the remainder rolled into the next page and produced a bimodal page-size pattern (e.g. 128B values, 64KB budget, 1024-row batches: 968 / 540 / 540 ... values per page). Return the boundary value's index + 1 instead, so the mini-batch crosses the threshold by exactly one value and the caller's page-flush check trips immediately, with no leftover sliver carried into the next page. The worst-case overshoot per page is one value's encoded size, which already matched the previous behavior whenever a single value alone exceeded the budget (the dropped .max(1) floor). Reported by Ed Seidel in apache#9972 review. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The mini-batch byte-budget walk now includes the value that crosses the budget, so the dictionary in the spill sub-test fills at 126 rows (1008 bytes) instead of 125 rows (1000 bytes), and the downstream plain page picks up 1254 rows / 10032 bytes instead of 1255 / 10040. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…dSizeBinary page sizing Adds two `LayoutTest` cases to `arrow_writer_layout.rs` that exercise byte-budget page-sizing paths introduced in apache#9972 through the real `ArrowWriter` user path: - `test_dictionary`: an arrow `DictionaryArray<Int32, Utf8>` input, which drives the dictionary-input arm of `ByteArrayEncoder::count_values_within_byte_budget_gather` (`DataType::Dictionary(_, _) => indices.len()`). Previously uncovered. - `test_fixed_size_binary`: a non-dictionary `FixedSizeBinary` column, which the arrow writer routes through the generic `ColumnValueEncoderImpl<FixedLenByteArrayType>`. Covers the FLBA branch of `plain_encoded_byte_size` and the variable-width scan in `count_values_within_byte_budget` via the arrow path (only the raw column-writer test covered it before). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
adriangb
force-pushed
the
parquet-page-size-mid-batch
branch
from
June 3, 2026 20:35
4769588 to
bcdb878
Compare
adriangb
marked this pull request as ready for review
June 3, 2026 20:35
adriangb
merged commit Jun 3, 2026
7d305a6
into
pydantic:parquet-page-size-mid-batch
15 of 18 checks passed
adriangb
added a commit
that referenced
this pull request
Jun 3, 2026
Fixes the CI failures from #16 (cargo fmt + clippy unnecessary_cast in arrow_writer_layout.rs) and applies that PR's follow-up recommendation: remove the unreachable count_within_budget_accessor free function and its `_ =>` accessor dispatch arm in ByteArrayEncoder. The arm was dead code: every byte-array type ByteArrayEncoder is built for has an explicit match arm, and a Dictionary(value = FixedSizeBinary) column hits the Dictionary(_, _) arm (its values.data_type() is Dictionary). A bare FixedSizeBinary column is routed to the generic column writer, never this encoder. The arm is now an unreachable! with that rationale, and the inaccurate "FixedSizeBinary falls through to the per-value walk" comment is dropped. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
adriangb
added a commit
that referenced
this pull request
Jun 3, 2026
Fixes the CI failures from #16 (cargo fmt + clippy unnecessary_cast in arrow_writer_layout.rs) and applies that PR's follow-up recommendation: remove the unreachable count_within_budget_accessor free function and its `_ =>` accessor dispatch arm in ByteArrayEncoder. The arm was dead code: every byte-array type ByteArrayEncoder is built for has an explicit match arm, and a Dictionary(value = FixedSizeBinary) column hits the Dictionary(_, _) arm (its values.data_type() is Dictionary). A bare FixedSizeBinary column is routed to the generic column writer, never this encoder. The arm is now an unreachable! with that rationale, and the inaccurate "FixedSizeBinary falls through to the per-value walk" comment is dropped. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Follow-up to apache#9972 — adds
LayoutTestcoverage for two byte-budget page-sizing paths that PR introduces, and reports one path that is unreachable dead code.This is a draft, stacked on top of apache#9972 (
pydantic:parquet-page-size-mid-batch). The tests are calibrated to the page-splitting behavior apache#9972 adds, so they must land on top of it.What changes are included in this PR?
Two new
LayoutTestcases inparquet/tests/arrow_writer_layout.rs, both driven through the realArrowWriteruser path:test_dictionary— an arrowDictionaryArray<Int32, Utf8>input. This exercises the dictionary-input arm ofByteArrayEncoder::count_values_within_byte_budget_gather(DataType::Dictionary(_, _) => indices.len()), which keeps an already-dictionary-encoded arrow column on the batched fast path. That arm was previously uncovered; with this testcargo llvm-covshows it hit (byte_array.rs:556).test_fixed_size_binary— a non-dictionaryFixedSizeBinarycolumn. The arrow writer routes this through the genericColumnValueEncoderImpl<FixedLenByteArrayType>(notByteArrayEncoder), so it covers the FLBA branch ofplain_encoded_byte_sizeand the variable-width scan incount_values_within_byte_budgetvia the arrow path — previously only the raw column-writer test reached that code.Note: one path could not be turned into a LayoutTest (it's dead code)
The original review flagged
count_within_budget_accessor(byte_array.rs:691–705) and its_ =>dispatch arm as uncovered. Investigating with coverage shows it is unreachable throughArrowWriter, not merely untested:get_arrow_column_writerroutes a bareFixedSizeBinarycolumn to the generic column writer, neverByteArrayEncoder.Dictionary(value = FixedSizeBinary)column does useByteArrayEncoder, but itsvalues.data_type()isDictionary, so it hits theDictionary(_, _)arm — not the accessor._ =>arm; every typeByteArrayEncoderis constructed for has an explicit match arm._ =>arm callsdowncast_op!, which has noFixedSizeBinarytop-level arm, so if it ever were reached it would hitunreachable!.A
FixedSizeBinaryLayoutTest confirms this: it passes through the encoder path (encoder.rsFLBA branch hit) whilecount_within_budget_accessorstays atMISS(0).Recommendation for apache#9972: remove
count_within_budget_accessor, the_ =>accessor arm, and the "FixedSizeBinary falls through to the per-value walk" comment (which is inaccurate), rather than trying to add a test for it. This also trims the diff.Are there any user-facing changes?
No — tests only.