This is an automated email from the ASF dual-hosted git repository.
etseidl pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 2541f87dea fix(parquet): cut data page byte-budget mini-batches on
exact value counts (#10554)
2541f87dea is described below
commit 2541f87deacb3a88ec12ccaadbc2759941e125af
Author: Adrian Garcia Badaracco <[email protected]>
AuthorDate: Fri Sep 4 07:29:15 2026 -0700
fix(parquet): cut data page byte-budget mini-batches on exact value counts
(#10554)
- closes https://github.com/apache/arrow-rs/issues/10538
> [!IMPORTANT]
> **Stacked on #10505 and #10745 — do not merge first.** The commits
below the top one belong to those PRs.
> Review only this PR's own commit:
[`cff2ca07fe...bd66b9ed68`](https://github.com/pydantic/arrow-rs/compare/cff2ca07fe875f38270c4db6cceb4a01d91ed51c...bd66b9ed68933cc4be79bb818299d294c97c74b6).
> Once #10505 and #10745 merge I'll rebase and this PR's diff becomes
clean on its own.
## The problem
`byte_budget_sub_batch_size` asks the encoder how many *values* fit in a
page byte budget, then converts that to a *level* count using the
chunk-wide level:value ratio, rounded up:
```rust
(values_per_subbatch * chunk_size).div_ceil(vals_in_chunk).max(1)
```
For a chunk with no nulls this is exact. With one null in 17 levels it
gives `ceil(17/16) == 2`, and `write_granular_chunk` slices the chunk
into uniform two-level windows — most of which carry **two** values,
i.e. twice what the budget allowed. The mechanism predates #10505; it
dates to #9972.
Where the encoding compresses a value against its predecessor, that
round-up costs whole values of output. 128 values of 2 MiB at one null
in 16:
| | file size |
| --- | --- |
| ratio-scaled windows | 16.78 MB |
| value-exact windows | **2.10 MB** |
At 8 MiB values it is 64 MiB against 8 MiB, and the acceptance case in
#10538 — 16 identical 64 KiB values with one null — goes from five pages
storing ~5 values in full to one page storing ~1.
## The fix
Have the chunker return the value count it already computed, and let
`write_granular_chunk` end a window by walking definition levels until
it has covered that many values. No ratio, no rounding.
Then apply it only where it changes the bytes written, because it is not
free — value-exact windows roughly double the mini-batch count on a
nullable column. Two conditions, both required:
**The budget must be the data page budget.** That one is a constant
`data_page_size_limit`, so a one-value budget means the value itself
overflows a page. The dictionary page budget is the limit *minus what
the dictionary already holds*, so it shrinks toward zero as the
dictionary fills and reaches a one-value budget on perfectly ordinary
values; cutting exactly there measured +13.0% on `string/default` and
+8.3% on `string/parquet_2`.
**The encoding must compress against the previous value.** `PLAIN` and
`DELTA_LENGTH_BYTE_ARRAY` store a value identically wherever it lands,
so value-exact windows leave their output byte for byte the same while
doubling the page count — measured +27.6% on a nullable column for no
reduction in output at all. The `compresses_against_previous_value` flag
#10505 added already marks exactly the right set.
## What the other paths give up
A ratio-scaled window spans `ceil(values × levels / values_in_chunk)`
levels. Where one value already fills the budget that covers **at most
two values**, whatever the null density, and exactly one wherever the
ratio is a whole number. Measured against a 1 MiB limit:
| values | nulls | encoding | max page, ratio | max page, value-exact |
floor |
| --- | --- | --- | --- | --- | --- |
| 2 MiB | 1-in-16 | `PLAIN` | 4.00× | 2.00× | 2.00× |
| 2 MiB | 1-in-4 | `PLAIN` | 4.00× | 2.00× | 2.00× |
| 2 MiB | 1-in-2 | `PLAIN` | 2.00× | 2.00× | 2.00× |
| 8 MiB | 1-in-16 | `PLAIN` | 16.00× | 8.00× | 8.00× |
The floor is what a page must hold: one value. So the concession is a
factor of two above an unavoidable minimum, it does not vary with null
density, and — the property #9972 exists for — it does not scale with
`write_batch_size`. Before #9972 a page took a whole mini-batch: 1024 ×
2 MiB, roughly 2000× the limit.
## Measurements
Base is #10505's head (measured at `fd806be5d3`; both branches have
since been rebased onto `main`, with the trees verified identical across
the rebase), so these isolate this PR. Local, run base → branch → base
on an idle machine, with the two base passes as a per-benchmark noise
floor. Benchmarks are the ones added in #10561.
| benchmark | before the encoding gate | after | noise |
| --- | --- | --- | --- |
| `..._nullable/plain` | +27.6% | **−0.8%** | 0.4% |
| `..._nullable_trailing/delta_byte_array` | −27.8% | **−27.5%** | 0.8%
|
| `..._nullable/delta_byte_array` | +7.6% | +7.4% | 0.0% |
| `..._nullable_dense/delta_byte_array` | −0.2% | +0.6% | 1.0% |
| `large_string_distinct_nullable/delta_byte_array` | +3.1% | +11.3% |
2.0% |
| `medium_string_shared_prefix_nullable/delta_byte_array` | +1.3% |
+1.7% | 1.2% |
Two costs remain, both on `DELTA_BYTE_ARRAY` where the byte win does not
materialise:
- **+7.4%** where output was already close to deduplicated (16.78 MB →
2.10 MB is still a 8× reduction, so this one pays for itself).
- **+11.3%** where the values share no prefix at all, and the file is
251.670105 MB against 251.670637 MB — no reduction. The writer cannot
know in advance whether values will share prefixes, so this is the
premium for the 8× win when they do.
Verified byte-identical output — same length, same hash — between base
and this PR for dictionary-encoded nullable columns (four shapes) and
for repeated columns (both encodings), confirming those paths are
untouched rather than merely unchanged in aggregate.
## Scope
Repeated columns are unchanged: records cannot span pages, so a record
holding several over-limit values still exceeds the budget. That is
inherent to the format.
## Tests
- `test_column_writer_delta_byte_array_nullable_shared_prefix_dedup` —
re-pinned from `[2, 2, 2, 2, 9]` to `[17]` and renamed, the layout
#10505 left a marker for.
- `test_column_writer_caps_page_size_with_sparse_nulls` — pins two
values per page under `PLAIN`, so it fails both if the encoding gate is
dropped (pages would hold one) and if the bound is lost (they would hold
many).
Full `parquet` suite green (1312 lib + integration), `fmt` and `clippy
-D warnings` clean.
## Note on #10505
Its `bool_to_int_with_if` trips `cargo clippy -- -D warnings`, which CI
runs — worth fixing on that branch too. Corrected here as part of
editing that test.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude <[email protected]>
---
parquet/src/column/writer/byte_budget_chunker.rs | 129 ++++++++---
parquet/src/column/writer/mod.rs | 259 ++++++++++++++++-------
2 files changed, 283 insertions(+), 105 deletions(-)
diff --git a/parquet/src/column/writer/byte_budget_chunker.rs
b/parquet/src/column/writer/byte_budget_chunker.rs
index 0d8eeb85ab..6e5917e11a 100644
--- a/parquet/src/column/writer/byte_budget_chunker.rs
+++ b/parquet/src/column/writer/byte_budget_chunker.rs
@@ -23,6 +23,47 @@ use crate::column::writer::encoder::ColumnValueEncoder;
use crate::file::properties::ResolvedColumnProperties;
use crate::schema::types::ColumnDescriptor;
+/// How [`write_granular_chunk`] should cut mini-batch windows in one chunk.
+///
+/// Cutting on exact value counts is the precise option and the expensive one:
+/// it roughly doubles the number of mini-batches on a nullable column, because
+/// the level:value ratio no longer rounds a window up to cover a second value.
+/// It is used only where that precision buys something.
+///
+/// [`write_granular_chunk`]: super::GenericColumnWriter::write_granular_chunk
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum SubBatchStrategy {
+ /// Cut after exactly this many values, walking definition levels to find
+ /// the boundary.
+ ///
+ /// Used against the data page budget for encodings that compress a value
+ /// against its predecessor, where the round-up below costs whole values of
+ /// output: 16 MiB rather than 2 MiB for 128 values at one null in 16. See
+ /// [#10538] for the measurements.
+ ///
+ /// [#10538]: https://github.com/apache/arrow-rs/issues/10538
+ Values(usize),
+ /// Cut after this many levels: the value budget scaled by the chunk's
+ /// level:value ratio, rounded up.
+ ///
+ /// Used everywhere else, because everywhere else the round-up changes no
+ /// bytes — the *dictionary page* budget shrinks toward zero as the
+ /// dictionary fills, so a one-value budget is routine there for ordinary
+ /// values, and `PLAIN` and `DELTA_LENGTH_BYTE_ARRAY` values cost the same
+ /// wherever they land — while value-exact windows cost throughput
+ /// ([#10554]).
+ ///
+ /// The round-up is bounded, which is what makes it an acceptable price. A
+ /// window spans `ceil(values * levels / values_in_chunk)` levels, so where
+ /// one value already fills the budget it covers at most two values,
+ /// whatever the null density. The page bound [#9972] added therefore still
+ /// holds; it is two values per page rather than one.
+ ///
+ /// [#9972]: https://github.com/apache/arrow-rs/pull/9972
+ /// [#10554]: https://github.com/apache/arrow-rs/pull/10554
+ Levels(usize),
+}
+
/// Picks byte-budget-aware mini-batch sizes for one column.
///
/// The parquet column writer checks the data page byte limit only *after*
@@ -88,12 +129,15 @@ impl ByteBudgetChunker {
}
}
- /// Decide how many levels at the start of a chunk belong in one
+ /// Decide how many *values* at the start of a chunk belong in one
/// mini-batch, so the mini-batch cannot overflow whichever page is
/// currently accumulating value bytes: the data page when plain-encoding,
- /// or the *dictionary* page while dictionary-encoding. A returned value
- /// smaller than `chunk_size` triggers granular sub-batching in
- /// `write_batch_internal`.
+ /// or the *dictionary* page while dictionary-encoding.
+ ///
+ /// `None` means the whole chunk fits in a single mini-batch — the common
+ /// case. `Some(_)` triggers granular sub-batching in
+ /// `write_batch_internal`; see [`SubBatchStrategy`] for why the data page
and
+ /// dictionary page budgets get different windowing.
///
/// While dictionary-encoding, the data page holds only small RLE indices,
/// but the dictionary page accumulates the distinct values themselves —
@@ -101,14 +145,14 @@ impl ByteBudgetChunker {
/// mini-batch. The per-mini-batch dictionary spill check would otherwise
/// let one mini-batch of large values balloon the dictionary page.
///
- /// Returns `chunk_size` immediately (no value inspection) when the chunk
- /// is empty, or when the column is a fixed-width type whose mini-batches
+ /// Returns `None` immediately (no value inspection) when the chunk is
+ /// empty, or when the column is a fixed-width type whose mini-batches
/// statically cannot overshoot the relevant page.
///
/// `#[inline]`: this is a tiny per-chunk dispatcher; the actual byte
- /// inspection lives in the out-of-line `byte_budget_sub_batch_size`.
+ /// inspection lives in the out-of-line `byte_budget_sub_batch`.
#[inline]
- pub(crate) fn pick_sub_batch_size<E: ColumnValueEncoder>(
+ pub(crate) fn pick_sub_batch<E: ColumnValueEncoder>(
&self,
encoder: &E,
values: &E::Values,
@@ -116,33 +160,40 @@ impl ByteBudgetChunker {
chunk_def: LevelDataRef<'_>,
values_offset: usize,
chunk_size: usize,
- ) -> usize {
+ ) -> Option<SubBatchStrategy> {
if chunk_size == 0 {
- return chunk_size;
+ return None;
}
- let budget = if encoder.has_dictionary() {
+ // The second element selects the windowing; see [`SubBatchStrategy`].
Only a
+ // constant data page budget under an encoding that compresses against
+ // the previous value is worth cutting exactly.
+ let (budget, value_exact) = if encoder.has_dictionary() {
if self.static_dict_always_fits {
- return chunk_size;
+ return None;
}
// Bound the mini-batch by the dictionary page's *remaining*
// budget (it accumulates across mini-batches until it spills).
- match encoder.estimated_dict_page_size() {
- Some(used) => self.dict_page_byte_limit.saturating_sub(used),
- None => return chunk_size,
- }
+ // An encoder that cannot size its dictionary page (`?`) leaves
+ // the chunk unsplit.
+ let used = encoder.estimated_dict_page_size()?;
+ (self.dict_page_byte_limit.saturating_sub(used), false)
} else {
if self.static_always_fits {
- return chunk_size;
+ return None;
}
- self.page_byte_limit
+ (
+ self.page_byte_limit,
+ encoder.compresses_against_previous_value(),
+ )
};
- self.byte_budget_sub_batch_size::<E>(
+ self.byte_budget_sub_batch::<E>(
values,
value_indices,
chunk_def,
values_offset,
chunk_size,
budget,
+ value_exact,
)
}
@@ -152,7 +203,8 @@ impl ByteBudgetChunker {
/// `#[inline(never)]` keeps this slow path out of the hot
/// `write_batch_internal` loop; numeric and bool columns never reach it.
#[inline(never)]
- fn byte_budget_sub_batch_size<E: ColumnValueEncoder>(
+ #[expect(clippy::too_many_arguments)]
+ fn byte_budget_sub_batch<E: ColumnValueEncoder>(
&self,
values: &E::Values,
value_indices: Option<&[usize]>,
@@ -160,14 +212,15 @@ impl ByteBudgetChunker {
values_offset: usize,
chunk_size: usize,
budget: usize,
- ) -> usize {
+ value_exact: bool,
+ ) -> Option<SubBatchStrategy> {
// How many of this chunk's levels carry an actual value. For a
// non-nullable, unrepeated column every level is a value, so
// `value_count` is O(1) (`Absent`/`Uniform` def levels); only
// nullable or nested columns pay the O(chunk_size) def-level scan.
let vals_in_chunk = chunk_def.value_count(chunk_size,
self.max_def_level);
if vals_in_chunk == 0 {
- return chunk_size;
+ return None;
}
// Ask the encoder how many of the next values fit in one page byte
// budget. Dispatch on whether the caller supplied gather indices;
@@ -184,20 +237,28 @@ impl ByteBudgetChunker {
}
};
match fit {
- None => chunk_size,
+ // The encoder cannot size these values; write the chunk whole.
+ None => None,
+ // All of the chunk's values fit in the budget — no sub-batching.
+ Some(values_per_subbatch) if values_per_subbatch >= vals_in_chunk
=> None,
Some(values_per_subbatch) => {
- // Convert the value count back into a level count. For a
- // non-nullable column this is a no-op; for nullable/nested
- // columns scale by the chunk's observed value-to-level
- // ratio.
- let levels_per_subbatch = if vals_in_chunk == chunk_size {
- values_per_subbatch
+ // `count_values_within_byte_budget` never reports zero, but a
+ // zero-wide window would not advance `write_granular_chunk`.
+ let values_per_subbatch = values_per_subbatch.max(1);
+ Some(if value_exact {
+ SubBatchStrategy::Values(values_per_subbatch)
} else {
- (values_per_subbatch * chunk_size)
- .div_ceil(vals_in_chunk)
- .max(1)
- };
- chunk_size.min(levels_per_subbatch.max(1))
+ // Scale to a level count. Inexact on nullable chunks, and
+ // deliberately so — see `SubBatchStrategy::Levels`.
+ let levels_per_subbatch = if vals_in_chunk == chunk_size {
+ values_per_subbatch
+ } else {
+ (values_per_subbatch * chunk_size)
+ .div_ceil(vals_in_chunk)
+ .max(1)
+ };
+
SubBatchStrategy::Levels(chunk_size.min(levels_per_subbatch))
+ })
}
}
}
diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs
index c0969c7a13..01c812051f 100644
--- a/parquet/src/column/writer/mod.rs
+++ b/parquet/src/column/writer/mod.rs
@@ -54,7 +54,7 @@ use crate::schema::types::{BasicTypeInfo, ColumnDescPtr,
ColumnDescriptor};
mod byte_budget_chunker;
pub(crate) mod encoder;
-use byte_budget_chunker::ByteBudgetChunker;
+use byte_budget_chunker::{ByteBudgetChunker, SubBatchStrategy};
macro_rules! downcast_writer {
($e:expr, $i:ident, $b:expr) => {
@@ -660,9 +660,9 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
// mini-batch (the common case — small or fixed-width values, no
// further page-size accounting needed), or must we fall back to
// byte-budget-aware sub-batching to keep a page from overshooting
- // `data_page_size_limit`? `pick_sub_batch_size` returns
- // `chunk_size` for the former.
- let sub_batch_size = chunker.pick_sub_batch_size(
+ // `data_page_size_limit`? `pick_sub_batch` returns `None` for
+ // the former, and otherwise how wide one mini-batch may be.
+ let sub_batch = chunker.pick_sub_batch(
&self.encoder,
values,
value_indices,
@@ -671,25 +671,28 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a,
E> {
chunk_size,
);
- if sub_batch_size >= chunk_size {
- values_offset += self.write_mini_batch(
- values,
- values_offset,
- value_indices,
- chunk_size,
- chunk_def,
- chunk_rep,
- )?;
- } else {
- values_offset += self.write_granular_chunk(
- values,
- values_offset,
- value_indices,
- chunk_size,
- chunk_def,
- chunk_rep,
- sub_batch_size,
- )?;
+ match sub_batch {
+ None => {
+ values_offset += self.write_mini_batch(
+ values,
+ values_offset,
+ value_indices,
+ chunk_size,
+ chunk_def,
+ chunk_rep,
+ )?;
+ }
+ Some(sub_batch) => {
+ values_offset += self.write_granular_chunk(
+ values,
+ values_offset,
+ value_indices,
+ chunk_size,
+ chunk_def,
+ chunk_rep,
+ sub_batch,
+ )?;
+ }
}
levels_offset = end_offset;
}
@@ -852,13 +855,19 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a,
E> {
})
}
- /// Writes a chunk in `sub_batch_size`-level sub-batches, checking the
- /// data page byte limit after each. This keeps the page size close to
+ /// Writes a chunk in sub-batches sized by `sub_batch`, checking the page
+ /// byte limit after each. This keeps the page size close to
/// `data_page_size_limit` instead of overshooting it by a whole chunk.
///
- /// For repeated/nested columns sub-batches step from one `rep == 0`
- /// boundary to the next so a record never spans data pages, matching
- /// the parquet format rule.
+ /// [`SubBatchStrategy::Values`] windows are cut on value boundaries,
+ /// walking definition levels to find where the value budget is used up;
+ /// [`SubBatchStrategy::Levels`] windows are a fixed level count. See
+ /// [`SubBatchStrategy`] for which budget gets which and why.
+ ///
+ /// For repeated/nested columns sub-batches then extend to the next
+ /// `rep == 0` boundary so a record never spans data pages, matching the
+ /// parquet format rule. A record holding several over-limit values
+ /// therefore still exceeds the budget; that is inherent to the format.
///
/// Returns the total number of values consumed across all sub-batches.
///
@@ -875,30 +884,40 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a,
E> {
chunk_size: usize,
chunk_def: LevelDataRef<'_>,
chunk_rep: LevelDataRef<'_>,
- sub_batch_size: usize,
+ sub_batch: SubBatchStrategy,
) -> Result<usize> {
- // The chunker always sizes a sub-batch to at least one level, so each
+ // The chunker always sizes a sub-batch to at least one value or one
+ // level, and a value-exact window spans at least one level, so each
// iteration below makes progress (`sub_end > sub_start`).
- debug_assert!(sub_batch_size >= 1, "chunker must size at least one
level");
+ debug_assert!(
+ matches!(sub_batch, SubBatchStrategy::Values(n) |
SubBatchStrategy::Levels(n) if n >= 1),
+ "chunker must size at least one value or level"
+ );
+ let max_def_level = self.descr.max_def_level();
let mut values_consumed = 0;
let mut sub_start = 0;
while sub_start < chunk_size {
+ let window_end = match sub_batch {
+ SubBatchStrategy::Values(n) => {
+ Self::window_end_for_values(chunk_def, chunk_size,
max_def_level, sub_start, n)
+ }
+ SubBatchStrategy::Levels(n) => (sub_start + n).min(chunk_size),
+ };
let sub_end = match chunk_rep {
LevelDataRef::Materialized(levels) => {
- // Pack up to `sub_batch_size` levels per mini-batch, then
- // extend to the next record boundary (rep == 0) so a
- // record never spans data pages. Packing whole records
- // rather than stepping one record at a time avoids
- // calling `write_mini_batch` per record: records average
- // only a handful of levels, so a record-at-a-time step
- // would issue many more mini-batches than necessary.
- let mut e = (sub_start + sub_batch_size).min(chunk_size);
+ // Extend the window to the next record boundary
+ // (rep == 0) so a record never spans data pages. Packing
+ // whole records rather than stepping one record at a time
+ // avoids calling `write_mini_batch` per record: records
+ // average only a handful of levels, so a record-at-a-time
+ // step would issue many more mini-batches than necessary.
+ let mut e = window_end;
while e < chunk_size && levels[e] != 0 {
e += 1;
}
e
}
- _ => (sub_start + sub_batch_size).min(chunk_size),
+ _ => window_end,
};
let sub_len = sub_end - sub_start;
let written = self.write_mini_batch(
@@ -915,6 +934,54 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E>
{
Ok(values_consumed)
}
+ /// Index one past the last level of a sub-batch window that starts at
+ /// `start` and covers at most `max_values` values, clamped to
+ /// `chunk_size`.
+ ///
+ /// Nulls trailing the last value are left to the next window, so a window
+ /// ends immediately after the value that exhausts its budget. The window
+ /// always spans at least one level, so callers make progress: `start` is
+ /// below `chunk_size` and `max_values` is at least one.
+ ///
+ /// Only nullable and nested columns pay the level walk. It runs solely in
+ /// the granular path, whose values are by definition large enough to
+ /// overflow a page budget, so touching each of the chunk's levels once is
+ /// noise next to writing those values — and that path already makes full
+ /// def-level and rep-level passes.
+ fn window_end_for_values(
+ chunk_def: LevelDataRef<'_>,
+ chunk_size: usize,
+ max_def_level: i16,
+ start: usize,
+ max_values: usize,
+ ) -> usize {
+ match chunk_def {
+ // `max_def_level == 0`: every level is a value.
+ LevelDataRef::Absent => (start + max_values).min(chunk_size),
+ LevelDataRef::Uniform { value, .. } => {
+ if value == max_def_level {
+ (start + max_values).min(chunk_size)
+ } else {
+ // Uniformly below max def: the chunk holds no values at
+ // all, so no window boundary can help. (The chunker
+ // returns `None` for such a chunk, so this is defensive.)
+ chunk_size
+ }
+ }
+ LevelDataRef::Materialized(levels) => {
+ let mut seen = 0;
+ let mut end = start;
+ while end < chunk_size && seen < max_values {
+ if levels[end] == max_def_level {
+ seen += 1;
+ }
+ end += 1;
+ }
+ end
+ }
+ }
+ }
+
/// Creates a new streaming level encoder appropriate for the writer
version.
fn create_level_encoder(max_level: i16, props: &WriterProperties) ->
LevelEncoder {
match props.writer_version() {
@@ -1095,12 +1162,13 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a,
E> {
/// `PLAIN` and `DELTA_LENGTH_BYTE_ARRAY` keep their tighter one-value page
/// bound.
///
- /// Known limitation: the caller's trigger keys on a page-opening
- /// mini-batch holding exactly one value. Nulls in a chunk make the
- /// byte-budget chunker emit multi-level mini-batches, so on nullable
- /// columns pages that open with a two-value mini-batch miss the
- /// exemption and dedup is only partial; see
- ///
`test_column_writer_delta_byte_array_nullable_shared_prefix_partial_dedup`.
+ /// The caller's trigger keys on a page-opening mini-batch holding exactly
+ /// one value. `write_granular_chunk` cuts windows after an exact value
+ /// count, so an over-limit value gets a single-value mini-batch whether
+ /// or not the chunk contains nulls; see
+ /// `test_column_writer_delta_byte_array_nullable_shared_prefix_dedup`.
+ /// The exception is a repeated column, where a record holding several
+ /// over-limit values cannot be split across pages at all.
#[cold]
fn set_page_size_exemption(&mut self) {
if !self.encoder.compresses_against_previous_value() {
@@ -3162,30 +3230,77 @@ mod tests {
}
#[test]
- fn
test_column_writer_delta_byte_array_nullable_shared_prefix_partial_dedup() {
- // Documents the *current* behavior of the first-value exemption on a
- // nullable column; this pins a known limitation, not an ideal.
+ fn test_column_writer_caps_page_size_with_sparse_nulls() {
+ // `PLAIN` keeps ratio-scaled windows, so a sparsely-null chunk puts
+ // *two* over-limit values on a page rather than one. That is the
+ // deliberate choice: value-exact windows would cut this to one, but
+ // `PLAIN` stores a value identically wherever it lands, so the output
+ // is byte for byte the same either way while the page count doubles.
//
- // The exemption fires when a page's first mini-batch contains exactly
- // one value. For a non-nullable column the byte-budget chunker gives
- // an over-limit value a one-level mini-batch, so that always holds.
- // One null in the chunk changes the level:value ratio to 17:16, the
- // chunker rounds up to two-level mini-batches, and a page whose first
- // mini-batch carries two values misses the exemption: it is cut after
- // those two values, and its first value is stored in full.
- //
- // The one mini-batch that pairs the null with a value has a single
- // value, so the page it opens does get the exemption and accumulates
- // every remaining suffix. The result for 16 identical values with a
- // null at index 8 is four two-value pages (each storing one value in
- // full), then one exempt page holding the rest:
- //
- // values per page: [2, 2, 2, 2, 9] (counts include the null level)
- // total bytes: ~5 full values, vs ~1 ideally and 16 for PLAIN
+ // What has to hold is that the bound stays a small constant and does
+ // not scale with `write_batch_size` — the failure #9972 fixed, where
+ // a page took a whole mini-batch. A window spans
+ // `ceil(values * levels / values_in_chunk)` levels, which covers at
+ // most two values however sparse the nulls are, so two is the whole
+ // exposure. Assert it exactly: this test fails if the encoding gate
+ // on value-exact windows is dropped (pages would hold one value) as
+ // well as if the bound is lost (they would hold many).
+ let value_size = 64 * 1024;
+ let page_byte_limit = 16 * 1024;
+ let num_values = 16;
+
+ let props = WriterProperties::builder()
+ .set_dictionary_enabled(false)
+ .set_encoding(Encoding::PLAIN)
+ .set_data_page_size_limit(page_byte_limit)
+ .set_statistics_enabled(EnabledStatistics::None)
+ .build();
+
+ let data: Vec<_> = (0..num_values)
+ .map(|_| ByteArray::from(vec![b'a'; value_size]))
+ .collect();
+ // 17 levels: a null at index 8, values everywhere else.
+ let def_levels: Vec<i16> = (0..num_values as i16 + 1)
+ .map(|i| i16::from(i != 8))
+ .collect();
+ let pages =
+ write_and_collect_pages::<ByteArrayType>(props, 1, 0, &data,
Some(&def_levels), None);
+
+ // At most two values' payload on any page, and never a whole
+ // mini-batch's worth.
+ let upper_bound = 2 * value_size + 64;
+ for (size, n_levels) in &pages.data_pages {
+ assert!(
+ *size <= upper_bound,
+ "page size {size} exceeds two-value bound ({upper_bound}B);
pages {:?}",
+ pages.data_pages,
+ );
+ assert!(
+ *n_levels <= 3,
+ "page holds {n_levels} levels, expected at most 3 (two values
+ a null); \
+ pages {:?}",
+ pages.data_pages,
+ );
+ }
+ // Two-level windows over 17 levels: eight pages carrying two levels
+ // and a ninth holding the remainder.
+ let num_levels: usize = num_values + 1;
+ assert_eq!(pages.data_pages.len(), num_levels.div_ceil(2));
+ }
+
+ #[test]
+ fn test_column_writer_delta_byte_array_nullable_shared_prefix_dedup() {
+ // A null in the chunk must not cost dedup. The first-value exemption
+ // fires when a page's opening mini-batch holds exactly one value, and
+ // `write_granular_chunk` cuts windows after an exact value count, so
+ // the over-limit value that opens the page gets a single-value
+ // mini-batch regardless of where nulls fall.
//
- // If the exemption trigger is ever keyed on values written to the
- // page (0 -> 1) instead of mini-batch shape, this test should fail
- // with fewer, larger pages — update it to pin the improved layout.
+ // This pinned `[2, 2, 2, 2, 9]` before #10538: the chunker scaled the
+ // one-value budget by the chunk's 17:16 level:value ratio and rounded
+ // up to two-level windows, so most pages opened with two values,
+ // missed the exemption, and stored their first value in full — ~5
+ // full values against ~1 here (`PLAIN` stores all 16).
let value_size = 64 * 1024;
let page_byte_limit = 16 * 1024;
let num_values = 16;
@@ -3208,13 +3323,15 @@ mod tests {
let pages =
write_and_collect_pages::<ByteArrayType>(props, 1, 0, &data,
Some(&def_levels), None);
+ // One page holding all 17 levels: the first value is exempt and every
+ // later value dedups down to a length pair.
let per_page_values: Vec<u32> = pages.data_pages.iter().map(|(_, n)|
*n).collect();
- assert_eq!(per_page_values, vec![2, 2, 2, 2, 9]);
+ assert_eq!(per_page_values, vec![num_values as u32 + 1]);
let total_bytes: usize = pages.data_pages.iter().map(|(size, _)|
size).sum();
assert!(
- total_bytes > 4 * value_size && total_bytes < 6 * value_size,
- "expected ~5 full values' worth of bytes (partial dedup), \
+ total_bytes < 2 * value_size,
+ "expected ~one full value's worth of bytes (full dedup), \
got {total_bytes}B across pages {:?}",
pages.data_pages,
);
