From ef40c4678440025c7a355ea6923fedc375b5291e Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Mon, 25 May 2026 12:44:25 -0500 Subject: [PATCH 01/15] test(parquet): add page-size regression tests for large values 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) --- parquet/src/arrow/arrow_writer/mod.rs | 239 +++++++++++++ parquet/src/column/writer/mod.rs | 494 ++++++++++++++++++++++++++ 2 files changed, 733 insertions(+) diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 79542caed9b7..7c3152d05266 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -4904,6 +4904,245 @@ mod tests { assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4); } + #[test] + fn test_arrow_writer_caps_page_size_for_large_strings() { + // End-to-end coverage for `ByteArrayEncoder::estimated_value_bytes_gather` + // and the offsets-buffer fast path in `estimate_byte_size_offsets`. + // + // The standard column-writer regression test covers the same + // logic for the non-arrow path; this one drives writes through + // `ArrowWriter`, which is what real users hit and which uses + // `write_gather` with `non_null_indices` (always contiguous for + // a non-null column). + let value_size = 64 * 1024; + let page_byte_limit = 16 * 1024; + let num_rows = 32; + + let schema = Arc::new(Schema::new(vec![Field::new( + "col", + ArrowDataType::Utf8, + false, + )])); + let strings: Vec = (0..num_rows).map(|_| "x".repeat(value_size)).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StringArray::from(strings)) as _], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(page_byte_limit) + .build(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + let data = Bytes::from(writer.into_inner().unwrap()); + + let mut metadata = ParquetMetaDataReader::new(); + metadata.try_parse(&data).unwrap(); + let metadata = metadata.finish().unwrap(); + let col_meta = metadata.row_group(0).column(0); + let mut reader = + SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap(); + + let mut data_pages = Vec::new(); + while let Some(page) = reader.get_next_page().unwrap() { + if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) { + data_pages.push(page.buffer().len()); + } + } + + assert!( + data_pages.len() >= num_rows / 2, + "expected at least {} data pages for {num_rows} large strings via ArrowWriter, got {} ({data_pages:?})", + num_rows / 2, + data_pages.len(), + ); + for size in &data_pages { + assert!( + *size <= value_size + 64, + "page size {size} exceeds one-value bound; pages {data_pages:?}" + ); + } + } + + #[test] + fn test_arrow_writer_granular_mode_roundtrip() { + // Granular mode subdivides chunks and writes more pages than + // `main`. Make sure the data we write back is bit-identical to + // what went in — page-count assertions elsewhere only prove + // pages were cut, not that the encoded data is correct. + // + // Mix value sizes so that the cumulative-byte-budget cutoff + // lands mid-chunk, exercising both batched and granular paths + // within the same `write_batch_internal` call. + let small = "tiny".to_string(); + let big = "x".repeat(64 * 1024); + let strings: Vec = (0..256) + .map(|i| { + if i % 16 == 0 { + big.clone() + } else { + small.clone() + } + }) + .collect(); + + let schema = Arc::new(Schema::new(vec![Field::new( + "col", + ArrowDataType::Utf8, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StringArray::from(strings.clone())) as _], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(16 * 1024) + .build(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + let data = Bytes::from(writer.into_inner().unwrap()); + + let mut reader = ParquetRecordBatchReader::try_new(data, 1024).unwrap(); + let read = reader.next().unwrap().unwrap(); + assert!(reader.next().is_none(), "expected one batch"); + let col = read + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(col.len(), strings.len()); + for (i, expected) in strings.iter().enumerate() { + assert_eq!( + col.value(i), + expected.as_str(), + "value mismatch at index {i}" + ); + } + } + + #[test] + fn test_arrow_writer_caps_page_size_for_large_string_view() { + // Coverage for the view/dict fallback in + // `ByteArrayEncoder::count_values_within_byte_budget_gather`. + // `Utf8View` arrays don't expose a single contiguous offsets + // buffer, so the gather method falls back to the per-value walk + // via `ArrayAccessor::value`. + use arrow_array::StringViewArray; + let value_size = 64 * 1024; + let page_byte_limit = 16 * 1024; + let num_rows = 32usize; + + let schema = Arc::new(Schema::new(vec![Field::new( + "col", + ArrowDataType::Utf8View, + false, + )])); + let strings: Vec = (0..num_rows).map(|_| "x".repeat(value_size)).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StringViewArray::from_iter_values( + strings.iter().map(|s| s.as_str()), + )) as _], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(page_byte_limit) + .build(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + let data = Bytes::from(writer.into_inner().unwrap()); + + let mut metadata = ParquetMetaDataReader::new(); + metadata.try_parse(&data).unwrap(); + let metadata = metadata.finish().unwrap(); + let col_meta = metadata.row_group(0).column(0); + let mut reader = + SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap(); + + let mut data_pages = Vec::new(); + while let Some(page) = reader.get_next_page().unwrap() { + if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) { + data_pages.push(page.buffer().len()); + } + } + assert!( + data_pages.len() >= num_rows / 2, + "expected pages bounded by byte budget for Utf8View, got {data_pages:?}" + ); + for size in &data_pages { + assert!( + *size <= value_size + 64, + "page size {size} exceeds one-value bound; pages {data_pages:?}" + ); + } + } + + #[test] + fn test_arrow_writer_all_null_string_column() { + // The `LevelDataRef::value_count` Uniform branch with + // `value != max_def` (entirely-null chunk) must return 0 so the + // sub-batch sizer short-circuits to batch mode without trying + // to estimate byte budgets for non-existent values. + let num_rows = 1024; + let schema = Arc::new(Schema::new(vec![Field::new( + "col", + ArrowDataType::Utf8, + true, + )])); + let nulls: Vec> = vec![None; num_rows]; + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StringArray::from(nulls)) as _], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(16 * 1024) + .build(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + let data = Bytes::from(writer.into_inner().unwrap()); + + // Re-parse the file: row group has one column, every row is + // null, all data pages report `num_rows / page_count` rows. + let mut metadata = ParquetMetaDataReader::new(); + metadata.try_parse(&data).unwrap(); + let metadata = metadata.finish().unwrap(); + let row_group = metadata.row_group(0); + let col_meta = row_group.column(0); + assert_eq!(row_group.num_rows() as usize, num_rows); + // Statistics record `null_count = num_rows` — proves every value + // was written as null. + if let Some(stats) = col_meta.statistics() { + assert_eq!( + stats.null_count_opt().unwrap_or(0) as usize, + num_rows, + "expected all-null column to report null_count = num_rows" + ); + } + + let mut reader = + SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap(); + let mut total_values = 0u32; + while let Some(page) = reader.get_next_page().unwrap() { + if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) { + total_values += page.num_values(); + } + } + assert_eq!( + total_values as usize, num_rows, + "expected every level position to be represented in some page" + ); + } + struct WriteBatchesShape { num_batches: usize, rows_per_batch: usize, diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index 4e53230bbf89..97ccf13a316b 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -2676,6 +2676,500 @@ mod tests { assert_eq!(other_values, vec![10]); } + #[test] + fn test_column_writer_caps_page_size_for_large_byte_array_values() { + // Regression: the post-write data page byte limit check only fires + // at mini-batch boundaries, so a 1024-row mini-batch of multi-MiB + // BYTE_ARRAY values used to buffer multiple GiB into a single page + // before the limit was even consulted. With the threshold-based + // granular mode this batch should split into ~one page per value. + let value_size = 64 * 1024; // 64 KiB per value + let page_byte_limit = 16 * 1024; // 16 KiB page limit + let num_rows = 64; + + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let props = Arc::new( + WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_data_page_size_limit(page_byte_limit) + // Default write_batch_size (1024) — without the fix this + // buffers the entire input into a single ~4 MiB page. + .build(), + ); + + let mut data = Vec::with_capacity(num_rows); + for i in 0..num_rows { + data.push(ByteArray::from(vec![i as u8; value_size])); + } + + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + writer.write_batch(&data, None, None).unwrap(); + let r = writer.close().unwrap(); + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let mut page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(), + ); + + let mut data_page_sizes = Vec::new(); + let mut data_page_value_counts = Vec::new(); + while let Some(page) = page_reader.get_next_page().unwrap() { + if matches!( + page.page_type(), + PageType::DATA_PAGE | PageType::DATA_PAGE_V2 + ) { + data_page_sizes.push(page.buffer().len()); + data_page_value_counts.push(page.num_values()); + } + } + + // Every value must end up somewhere. + assert_eq!( + data_page_value_counts.iter().sum::() as usize, + num_rows + ); + // Without the fix this assertion fired with one ~4 MiB page; the + // threshold splits the input so that no page holds more than a + // single oversized value's worth of bytes. + assert!( + data_page_sizes.len() >= num_rows / 2, + "expected pages to be cut close to one per value, got {} pages for sizes {:?}", + data_page_sizes.len(), + data_page_sizes, + ); + // Each page must be bounded by roughly one value's worth of bytes; + // parquet allows a single oversized value to occupy a page by + // itself but never lets us pile many of them together. + for size in &data_page_sizes { + assert!( + *size <= value_size + 64, + "page size {size} exceeds one-value bound ({}B) — pages {:?}", + value_size + 64, + data_page_sizes, + ); + } + } + + #[test] + fn test_column_writer_caps_page_size_for_large_values_in_list() { + // Coverage for the Materialized-rep branch of + // `write_granular_chunk`. The flat-column regression test + // exercises the per-level step; this exercises the + // record-by-record step used when rep levels are present. + // + // Column is `list` (max_def = 1, max_rep = 1) + // with 3 records of 3 large blobs each. The page byte limit is + // smaller than a single blob, so granular mode kicks in, and the + // Materialized-rep arm of `write_granular_chunk` steps from one + // `rep == 0` boundary to the next so a record never spans pages. + let value_size = 32 * 1024; + let page_byte_limit = 16 * 1024; + let values_per_record = 3; + let num_records = 3; + let num_values = values_per_record * num_records; + + // rep levels: 0, 1, 1, 0, 1, 1, 0, 1, 1 + let mut rep_levels = Vec::with_capacity(num_values); + for _ in 0..num_records { + rep_levels.push(0i16); + rep_levels.extend(std::iter::repeat_n(1i16, values_per_record - 1)); + } + let def_levels = vec![1i16; num_values]; + + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let props = Arc::new( + WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_data_page_size_limit(page_byte_limit) + .build(), + ); + + let mut data = Vec::with_capacity(num_values); + for i in 0..num_values { + data.push(ByteArray::from(vec![i as u8; value_size])); + } + + let mut writer = get_test_column_writer::(page_writer, 1, 1, props); + writer + .write_batch(&data, Some(&def_levels), Some(&rep_levels)) + .unwrap(); + let r = writer.close().unwrap(); + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let mut page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(), + ); + + let mut data_pages = Vec::new(); + while let Some(page) = page_reader.get_next_page().unwrap() { + if matches!( + page.page_type(), + PageType::DATA_PAGE | PageType::DATA_PAGE_V2 + ) { + data_pages.push((page.buffer().len(), page.num_values())); + } + } + + // The Materialized-rep arm groups levels by record, and each + // record's bytes blow the page byte limit on its own, so we get + // exactly one page per record. + assert_eq!( + data_pages.len(), + num_records, + "expected one data page per record, got {data_pages:?}" + ); + for (bytes, n_values) in &data_pages { + assert_eq!( + *n_values as usize, values_per_record, + "each page must hold a whole record's leaves, got {data_pages:?}" + ); + // Each page is one full record (its leaves cannot be split), + // so allow up to `values_per_record` blobs of payload plus a + // small fudge for level encoding overhead. + let upper_bound = values_per_record * (value_size + 16); + assert!( + *bytes <= upper_bound, + "page size {bytes} exceeds whole-record bound ({upper_bound}); pages {data_pages:?}" + ); + } + } + + #[test] + fn test_column_writer_caps_page_size_with_nullable_large_values() { + // Coverage for `LevelDataRef::value_count` on Materialized def + // levels: a nullable column with mixed nulls and large values. + // `value_count` must return the actual non-null count so the + // byte estimate reflects bytes that will actually be written, + // not the level count. + let value_size = 32 * 1024; + let page_byte_limit = 16 * 1024; + let num_levels = 32; + + // Alternating null / non-null: 16 nulls and 16 values. + let def_levels: Vec = (0..num_levels as i16).map(|i| i % 2).collect(); + let num_values = def_levels.iter().filter(|&&d| d == 1).count(); + + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let props = Arc::new( + WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_data_page_size_limit(page_byte_limit) + .build(), + ); + + let mut data = Vec::with_capacity(num_values); + for i in 0..num_values { + data.push(ByteArray::from(vec![i as u8; value_size])); + } + + let mut writer = get_test_column_writer::(page_writer, 1, 0, props); + writer.write_batch(&data, Some(&def_levels), None).unwrap(); + let r = writer.close().unwrap(); + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let mut page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(), + ); + + let mut data_pages = Vec::new(); + while let Some(page) = page_reader.get_next_page().unwrap() { + if matches!( + page.page_type(), + PageType::DATA_PAGE | PageType::DATA_PAGE_V2 + ) { + data_pages.push(page.buffer().len()); + } + } + + // With 16 actual values of 32 KiB each and a 16 KiB page limit, + // every non-null value should get its own page (plus possibly + // adjacent nulls). At minimum, the number of pages must be + // roughly the value count, not 1 (which is what `main` produced). + assert!( + data_pages.len() >= num_values / 2, + "expected at least {} pages for {num_values} large values, got {} pages: {data_pages:?}", + num_values / 2, + data_pages.len(), + ); + // No page contains more than ~one value's worth of payload bytes. + for size in &data_pages { + assert!( + *size <= value_size + 64, + "page size {size} exceeds one-value bound; pages {data_pages:?}" + ); + } + } + + #[test] + fn test_column_writer_dict_enabled_large_values_post_spill() { + // While dictionary encoding is active, `has_dictionary()` short- + // circuits `estimated_value_bytes` — the byte estimate is plain- + // encoded size but dict-encoded pages only store small RLE + // indices, so we'd otherwise shrink pages spuriously. Once the + // dictionary spills (each value is large + unique), plain + // encoding takes over and the byte-budget sub-batch kicks in. + // + // This test makes sure the writer survives that transition and + // produces bounded pages thereafter. + let value_size = 64 * 1024; + let page_byte_limit = 16 * 1024; + let num_rows = 32; + + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let props = Arc::new( + WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(true) + // Force a small dict so it spills quickly even though + // each value here is unique. + .set_dictionary_page_size_limit(1024) + .set_data_page_size_limit(page_byte_limit) + // Small mini-batches so dict fallback happens part-way + // through the input, leaving subsequent mini-batches to + // exercise the post-spill plain-encoding path that the + // page-size fix actually targets. + .set_write_batch_size(4) + .build(), + ); + + let mut data = Vec::with_capacity(num_rows); + for i in 0..num_rows { + data.push(ByteArray::from(vec![i as u8; value_size])); + } + + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + writer.write_batch(&data, None, None).unwrap(); + let r = writer.close().unwrap(); + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let mut page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(), + ); + + let mut data_pages = Vec::new(); + while let Some(page) = page_reader.get_next_page().unwrap() { + if matches!( + page.page_type(), + PageType::DATA_PAGE | PageType::DATA_PAGE_V2 + ) { + data_pages.push(page.buffer().len()); + } + } + + // After spill, plain encoding writes one ~64 KiB value per page. + // Without the fix, post-spill writes still buffered all 32 + // values into a single ~2 MiB page. + assert!( + data_pages.len() >= num_rows / 2, + "expected >= {} data pages after dict spill, got {} ({data_pages:?})", + num_rows / 2, + data_pages.len(), + ); + for size in &data_pages { + assert!( + *size <= value_size + 64, + "page size {size} exceeds one-value bound; pages {data_pages:?}" + ); + } + } + + #[test] + fn test_column_writer_caps_dictionary_page_size() { + // A column of large *distinct* values with dictionary encoding on: + // the dictionary page accumulates the values themselves, and its + // spill check runs only once per mini-batch. Without bounding the + // dictionary-encoding mini-batch, one `write_batch_size` mini-batch + // would intern `write_batch_size * value_size` bytes into the + // dictionary page before the check fires (~16 MiB here). The chunker + // must sub-batch the dictionary-encoding phase too. + let value_size = 8 * 1024; + let dict_page_limit = 64 * 1024; + let num_rows = 2048; + + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let props = Arc::new( + WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(true) + .set_dictionary_page_size_limit(dict_page_limit) + .build(), + ); + + let mut data = Vec::with_capacity(num_rows); + for i in 0..num_rows { + // each value distinct, so the dictionary cannot dedup them + let mut v = vec![0u8; value_size]; + v[..8].copy_from_slice(&(i as u64).to_le_bytes()); + data.push(ByteArray::from(v)); + } + + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + writer.write_batch(&data, None, None).unwrap(); + let r = writer.close().unwrap(); + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let mut page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(), + ); + + let mut dict_page_size = 0; + while let Some(page) = page_reader.get_next_page().unwrap() { + if page.page_type() == PageType::DICTIONARY_PAGE { + dict_page_size = dict_page_size.max(page.buffer().len()); + } + } + + assert!( + dict_page_size > 0, + "expected the column to dictionary-encode" + ); + // Bounded near the limit (~2x from the post-mini-batch check). Before + // the fix the dictionary page reached num_rows * value_size (~16 MiB, + // 256x the limit). + assert!( + dict_page_size <= 3 * dict_page_limit, + "dictionary page {dict_page_size} exceeds 3x the {dict_page_limit} limit", + ); + } + + #[test] + fn test_column_writer_caps_page_size_for_fixed_len_byte_array() { + // Coverage for `ParquetValueType::byte_size` override on + // `FixedLenByteArray`. With `type_length = 1`, each plain-encoded + // value is one byte, so a 4-byte page byte limit forces the + // sub-batch sizer to write ~4 values per page rather than one + // page for the whole batch. + let page_byte_limit = 4; + let num_values = 128; + + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let props = Arc::new( + WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_data_page_size_limit(page_byte_limit) + .build(), + ); + + let mut data = Vec::with_capacity(num_values); + for i in 0..num_values { + let mut fla = FixedLenByteArray::default(); + fla.set_data(Bytes::from(vec![i as u8])); + data.push(fla); + } + + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + writer.write_batch(&data, None, None).unwrap(); + let r = writer.close().unwrap(); + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let mut page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(), + ); + + let mut data_pages = Vec::new(); + while let Some(page) = page_reader.get_next_page().unwrap() { + if matches!( + page.page_type(), + PageType::DATA_PAGE | PageType::DATA_PAGE_V2 + ) { + data_pages.push(page.buffer().len()); + } + } + + // Without the fix this is a single 128-byte page; with the fix + // the byte budget caps each page at ~`page_byte_limit` bytes. + assert!( + data_pages.len() >= num_values / 8, + "expected pages capped by byte budget, got {data_pages:?}" + ); + for size in &data_pages { + assert!( + *size <= page_byte_limit * 4, + "page size {size} larger than expected; pages {data_pages:?}" + ); + } + } + #[test] fn test_bool_statistics() { let stats = statistics_roundtrip::(&[true, false, false, true]); From 17f7cc8b1793345abc4148609f1d3c0988991fbc Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Mon, 25 May 2026 12:45:28 -0500 Subject: [PATCH 02/15] fix(parquet): bound data page byte size for large variable-width values MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../src/column/writer/byte_budget_chunker.rs | 183 ++++++++++++++++++ parquet/src/column/writer/encoder.rs | 128 ++++++++++++ parquet/src/column/writer/mod.rs | 167 +++++++++++++++- 3 files changed, 472 insertions(+), 6 deletions(-) create mode 100644 parquet/src/column/writer/byte_budget_chunker.rs diff --git a/parquet/src/column/writer/byte_budget_chunker.rs b/parquet/src/column/writer/byte_budget_chunker.rs new file mode 100644 index 000000000000..465568b4caf6 --- /dev/null +++ b/parquet/src/column/writer/byte_budget_chunker.rs @@ -0,0 +1,183 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Decides how many levels of a chunk to write as one mini-batch so that +//! the resulting data page stays within `data_page_size_limit`. +//! +//! The parquet column writer checks the data page byte limit only *after* +//! each mini-batch finishes writing. Mini-batches are sized in rows +//! (`write_batch_size`, default 1024), so for BYTE_ARRAY columns whose +//! values are large (e.g. multi-MiB blobs) a single mini-batch can buffer +//! GiB into one page before the limit is consulted. +//! +//! This module isolates the per-chunk decision that prevents that: given a +//! chunk's level data and the input values, pick the largest `sub_batch_size` +//! such that one mini-batch fits in one page byte budget. For the +//! overwhelmingly common case (small or fixed-width values) the answer is +//! just `chunk_size` and the decision is O(1) on the column type — only +//! when the input might overflow does the chunker consult the encoder's +//! byte estimate. + +use crate::basic::Type; +use crate::column::writer::LevelDataRef; +use crate::column::writer::encoder::ColumnValueEncoder; +use crate::file::properties::WriterProperties; +use crate::schema::types::ColumnDescriptor; + +/// Picks byte-budget-aware mini-batch sizes for one column. +pub(crate) struct ByteBudgetChunker { + /// Configured data page byte limit for the column. + page_byte_limit: usize, + /// Max definition level of the column; a level equal to this marks a + /// present (non-null) leaf value. Used to count values per chunk. + max_def_level: i16, + /// `true` when no chunk of `base_batch_size` values can ever overflow + /// `page_byte_limit` regardless of input. Set once at column open from + /// the physical type's known per-value byte size; lets the per-chunk + /// decision short-circuit with no work for every numeric, bool, or + /// narrow `FIXED_LEN_BYTE_ARRAY` column. + static_always_fits: bool, +} + +impl ByteBudgetChunker { + #[inline] + pub(crate) fn new( + descr: &ColumnDescriptor, + props: &WriterProperties, + base_batch_size: usize, + ) -> Self { + let page_byte_limit = props.column_data_page_size_limit(descr.path()); + let static_bytes_per_value = match descr.physical_type() { + Type::BOOLEAN => Some(1), + Type::INT32 | Type::FLOAT => Some(std::mem::size_of::()), + Type::INT64 | Type::DOUBLE => Some(std::mem::size_of::()), + Type::INT96 => Some(12), + Type::FIXED_LEN_BYTE_ARRAY => Some(descr.type_length().max(0) as usize), + Type::BYTE_ARRAY => None, + }; + let static_always_fits = static_bytes_per_value + .map(|b| b.saturating_mul(base_batch_size) <= page_byte_limit) + .unwrap_or(false); + Self { + page_byte_limit, + max_def_level: descr.max_def_level(), + static_always_fits, + } + } + + /// Decide how many levels at the start of a chunk belong in one + /// mini-batch, so the mini-batch cannot overflow the data page that is + /// currently accumulating value bytes. A returned value smaller than + /// `chunk_size` triggers granular sub-batching in + /// `write_batch_internal`. + /// + /// Returns `chunk_size` immediately (no value inspection) when the chunk + /// is empty, when the column is dictionary-encoding (the data page then + /// holds only small RLE indices, so the plain-encoded byte estimate does + /// not apply), or when the column is a fixed-width type whose + /// mini-batches statically cannot overshoot the data page. + /// + /// `#[inline]`: this is a tiny per-chunk dispatcher; the actual byte + /// inspection lives in the out-of-line `byte_budget_sub_batch_size`. + #[inline] + pub(crate) fn pick_sub_batch_size( + &self, + encoder: &E, + values: &E::Values, + value_indices: Option<&[usize]>, + chunk_def: LevelDataRef<'_>, + values_offset: usize, + chunk_size: usize, + ) -> usize { + if chunk_size == 0 { + return chunk_size; + } + // While dictionary-encoding, the data page holds only small RLE + // indices — the plain-encoded byte estimate would spuriously shrink + // pages — so stay on the batched fast path. + if encoder.has_dictionary() { + return chunk_size; + } + if self.static_always_fits { + return chunk_size; + } + self.byte_budget_sub_batch_size::( + values, + value_indices, + chunk_def, + values_offset, + chunk_size, + self.page_byte_limit, + ) + } + + /// Inspect value sizes to decide how many of the chunk's values fit in + /// `budget` bytes (the data page remaining budget). + /// + /// `#[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( + &self, + values: &E::Values, + value_indices: Option<&[usize]>, + chunk_def: LevelDataRef<'_>, + values_offset: usize, + chunk_size: usize, + budget: usize, + ) -> usize { + // 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; + } + // Ask the encoder how many of the next values fit in one page byte + // budget. Dispatch on whether the caller supplied gather indices; + // this mirrors how `write_mini_batch` picks `write_gather` vs + // `write`. + let fit = match value_indices { + Some(idx) => { + let end = (values_offset + vals_in_chunk).min(idx.len()); + let start = values_offset.min(end); + E::count_values_within_byte_budget_gather(values, &idx[start..end], budget) + } + None => { + E::count_values_within_byte_budget(values, values_offset, vals_in_chunk, budget) + } + }; + match fit { + None => chunk_size, + 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 + } else { + (values_per_subbatch * chunk_size) + .div_ceil(vals_in_chunk) + .max(1) + }; + chunk_size.min(levels_per_subbatch.max(1)) + } + } + } +} diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index 2ea3376ae708..83c08e7423fa 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -90,6 +90,42 @@ pub trait ColumnValueEncoder { /// Write the values at the indexes in `indices` to this [`ColumnValueEncoder`] fn write_gather(&mut self, values: &Self::Values, indices: &[usize]) -> Result<()>; + /// Returns the largest `k` such that the first `k` values in + /// `values[offset..offset + len]` encode to at most `byte_budget` + /// bytes — i.e. how many values fit in a single page byte budget. + /// + /// Returns `len` if every value fits. Returns at least 1 if a single + /// value alone exceeds the budget, matching parquet's "at least one + /// value per data page" rule. + /// + /// `None` means "no cheap estimate available"; the caller stays on + /// the batched fast path and lets the post-write + /// `should_add_data_page` check handle bounding. + /// + /// Implementations should short-circuit aggressively: the typical + /// case is "everything fits, return `len`", and the next-most-common + /// case is "one wide value, return 1." The variable-width walk only + /// needs to be precise when the chunk is genuinely near the budget. + fn count_values_within_byte_budget( + _values: &Self::Values, + _offset: usize, + _len: usize, + _byte_budget: usize, + ) -> Option { + None + } + + /// As [`Self::count_values_within_byte_budget`] but using gather + /// `indices` rather than a contiguous range. Returns the number of + /// `indices` that fit, not the maximum index value. + fn count_values_within_byte_budget_gather( + _values: &Self::Values, + _indices: &[usize], + _byte_budget: usize, + ) -> Option { + None + } + /// Returns the number of buffered values fn num_values(&self) -> usize; @@ -247,6 +283,63 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { self.write_slice(&slice) } + fn count_values_within_byte_budget( + values: &[T::T], + offset: usize, + len: usize, + byte_budget: usize, + ) -> Option { + // Clamp so that a caller-supplied `len` that overruns the input + // (e.g. a level/value mismatch the encoder will reject later) + // returns an estimate instead of panicking here. + let end = (offset + len).min(values.len()); + let start = offset.min(end); + let n = end - start; + // Fixed-size physical types have a constant per-value byte cost, + // so the answer is one division — no need to walk the slice. + let phys = ::PHYSICAL_TYPE; + if phys != Type::BYTE_ARRAY && phys != Type::FIXED_LEN_BYTE_ARRAY { + let per = std::mem::size_of::().max(1); + let fits = (byte_budget / per).max(1); + return Some(fits.min(n)); + } + // Variable-width: scan, accumulate, exit at the first value that + // pushes us past the budget. This both bounds skewed + // distributions (one outlier among small values is caught when + // it lands, regardless of position) and short-circuits when an + // early value alone exceeds the budget. + let mut cum: usize = 0; + for (i, v) in values[start..end].iter().enumerate() { + cum = cum.saturating_add(plain_encoded_byte_size::(v)); + if cum > byte_budget { + return Some(i.max(1)); + } + } + Some(n) + } + + fn count_values_within_byte_budget_gather( + values: &[T::T], + indices: &[usize], + byte_budget: usize, + ) -> Option { + let phys = ::PHYSICAL_TYPE; + if phys != Type::BYTE_ARRAY && phys != Type::FIXED_LEN_BYTE_ARRAY { + let per = std::mem::size_of::().max(1); + let fits = (byte_budget / per).max(1); + return Some(fits.min(indices.len())); + } + let mut cum: usize = 0; + for (i, idx) in indices.iter().enumerate() { + let Some(v) = values.get(*idx) else { continue }; + cum = cum.saturating_add(plain_encoded_byte_size::(v)); + if cum > byte_budget { + return Some(i.max(1)); + } + } + Some(indices.len()) + } + fn num_values(&self) -> usize { self.num_values } @@ -411,3 +504,38 @@ where } } } + +/// Plain-encoded byte cost of a single value of type `T::T`. +/// +/// Derived from [`ParquetValueType::dict_encoding_size`] so we don't add a +/// parallel per-value-size hook to the trait. The components returned by +/// `dict_encoding_size` are `(per-value overhead, value-bytes)`. For +/// plain encoding the on-disk layout is: +/// +/// - `BYTE_ARRAY`: 4-byte length prefix + payload bytes = `overhead + bytes`. +/// - `FIXED_LEN_BYTE_ARRAY`: raw bytes only, taken from the type descriptor's +/// `type_length`. The value's own `dict_encoding_size` reports the length +/// prefix, which is irrelevant for plain FLBA encoding; the encoder passes +/// `type_length` directly. +/// - Everything else (numeric / bool): a constant per-value size; the caller +/// already short-circuits these via `mem::size_of::()` before +/// touching this function, so this branch is unreachable in practice and +/// we fall back to `overhead` defensively. +/// +/// See `dict_encoder.rs::push` (line ~52) for the matching dispatch. +/// +/// Placed at the end of the module deliberately. Inserting it above the +/// `ColumnValueEncoder` trait shifts the trait and `ColumnValueEncoderImpl` +/// within the compiled module enough to perturb downstream code placement, +/// which measurably regresses unrelated arrow-writer string benchmarks +/// (~5-9% on `string` / `string_and_binary_view`). Defining it last keeps +/// the hot encoder code at the offsets it has on `main`. +#[inline] +fn plain_encoded_byte_size(value: &T::T) -> usize { + let (overhead, bytes) = value.dict_encoding_size(); + match ::PHYSICAL_TYPE { + Type::BYTE_ARRAY => overhead + bytes, + Type::FIXED_LEN_BYTE_ARRAY => bytes, + _ => overhead, + } +} diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index 97ccf13a316b..52e2ea06a50c 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -49,8 +49,11 @@ use crate::file::properties::{ use crate::file::statistics::{Statistics, ValueStatistics}; use crate::schema::types::{ColumnDescPtr, ColumnDescriptor}; +mod byte_budget_chunker; pub(crate) mod encoder; +use byte_budget_chunker::ByteBudgetChunker; + macro_rules! downcast_writer { ($e:expr, $i:ident, $b:expr) => { match $e { @@ -374,6 +377,24 @@ impl<'a> LevelDataRef<'a> { Self::Uniform { value, .. } => Self::Uniform { value, count: len }, } } + + /// Count of positions in this slice that represent an actual value + /// (definition level equal to `max_def`). `Absent` means the column has + /// `max_def == 0` and every position is a value, so the implicit count + /// is the caller-supplied `total`. + pub(crate) fn value_count(self, total: usize, max_def: i16) -> usize { + match self { + Self::Absent => total, + Self::Materialized(values) => values.iter().filter(|&&d| d == max_def).count(), + Self::Uniform { value, count } => { + if value == max_def { + count + } else { + 0 + } + } + } + } } /// Typed column writer for a primitive column. @@ -545,6 +566,7 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { } else { self.props.write_batch_size() }; + let chunker = ByteBudgetChunker::new(&self.descr, &self.props, base_batch_size); while levels_offset < num_levels { let mut end_offset = num_levels.min(levels_offset + base_batch_size); @@ -555,14 +577,39 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { } } - values_offset += self.write_mini_batch( + let chunk_size = end_offset - levels_offset; + let chunk_def = def_levels.slice(levels_offset, chunk_size); + let chunk_rep = rep_levels.slice(levels_offset, chunk_size); + + let sub_batch_size = chunker.pick_sub_batch_size( + &self.encoder, values, - values_offset, value_indices, - end_offset - levels_offset, - def_levels.slice(levels_offset, end_offset - levels_offset), - rep_levels.slice(levels_offset, end_offset - levels_offset), - )?; + chunk_def, + values_offset, + 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, + )?; + } levels_offset = end_offset; } @@ -713,6 +760,78 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { }) } + /// Writes a chunk in sub-batches sized by the caller-supplied + /// `sub_batch_size` so the post-write data page byte limit check + /// fires before the chunk can grossly overshoot + /// `data_page_size_limit`. + /// + /// For flat (unrepeated) columns sub-batches contain up to + /// `sub_batch_size` levels each. 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. + /// + /// Returns the total number of values consumed across all sub-batches. + /// + /// `#[inline(never)]` keeps this slow path — only reached for + /// variable-width columns whose values need page splitting — out of + /// the hot `write_batch_internal` loop. + #[allow(clippy::too_many_arguments)] + #[inline(never)] + fn write_granular_chunk( + &mut self, + values: &E::Values, + values_offset: usize, + value_indices: Option<&[usize]>, + chunk_size: usize, + chunk_def: LevelDataRef<'_>, + chunk_rep: LevelDataRef<'_>, + sub_batch_size: usize, + ) -> Result { + let mut values_consumed = 0; + let mut sub_start = 0; + while sub_start < 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 + // emitting a `write_mini_batch` per record: records + // average only a handful of levels, so the + // record-at-a-time loop issued roughly `sub_batch_size`× + // more mini-batches than necessary. + let mut e = (sub_start + sub_batch_size).min(chunk_size); + while e < chunk_size && levels[e] != 0 { + e += 1; + } + // `sub_batch_size` can be 0 when the chunker sizes a + // sub-batch below one record; always make progress by + // consuming at least the first whole record. + if e == sub_start { + e = sub_start + 1; + while e < chunk_size && levels[e] != 0 { + e += 1; + } + } + e + } + _ => (sub_start + sub_batch_size).min(chunk_size), + }; + let sub_len = sub_end - sub_start; + let written = self.write_mini_batch( + values, + values_offset + values_consumed, + value_indices, + sub_len, + chunk_def.slice(sub_start, sub_len), + chunk_rep.slice(sub_start, sub_len), + )?; + values_consumed += written; + sub_start = sub_end; + } + Ok(values_consumed) + } + /// 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() { @@ -5211,6 +5330,42 @@ mod tests { } } + #[test] + fn test_level_data_ref_value_count() { + // `value_count` is what the byte-budget chunker uses to convert a + // chunk's level span into a leaf-value count. It must work for any + // column shape — flat, nullable, or nested — because the leaf + // values array is decoupled from the rep/def level stream. + let max_def = 2; + // Non-nullable / unrepeated: no def levels materialized — every + // level is a value. + assert_eq!(LevelDataRef::Absent.value_count(64, max_def), 64); + // Uniform run of present values, and of nulls. + assert_eq!( + LevelDataRef::Uniform { + value: max_def, + count: 40 + } + .value_count(40, max_def), + 40 + ); + assert_eq!( + LevelDataRef::Uniform { + value: max_def - 1, + count: 40 + } + .value_count(40, max_def), + 0 + ); + // Materialized def levels (nullable / nested): only levels equal to + // `max_def` are values; empty-list / null levels are not. + let levels = [2i16, 0, 2, 1, 2, 2, 0]; + assert_eq!( + LevelDataRef::Materialized(&levels).value_count(levels.len(), max_def), + 4 + ); + } + #[test] fn test_uniform_def_levels_all_null() { // All-null column: def_level=0 (null) for every slot, no values written. From 0ee44d96ea5da433fb36828ff43a5c8402de30f8 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Mon, 25 May 2026 12:45:38 -0500 Subject: [PATCH 03/15] fix(parquet): bound page size for arrow byte-array column writes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- parquet/src/arrow/arrow_writer/byte_array.rs | 223 ++++++++++++++++++- 1 file changed, 222 insertions(+), 1 deletion(-) diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index 9cb0718b4d84..c7e172f1a815 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -30,10 +30,12 @@ use crate::geospatial::statistics::GeospatialStatistics; use crate::schema::types::ColumnDescPtr; use crate::util::bit_util::num_required_bits; use crate::util::interner::{Interner, Storage}; +use arrow_array::types::ByteArrayType; use arrow_array::{ Array, ArrayAccessor, BinaryArray, BinaryViewArray, DictionaryArray, FixedSizeBinaryArray, - LargeBinaryArray, LargeStringArray, StringArray, StringViewArray, + GenericByteArray, LargeBinaryArray, LargeStringArray, StringArray, StringViewArray, }; +use arrow_buffer::{ArrowNativeType, Buffer}; use arrow_schema::DataType; macro_rules! downcast_dict_impl { @@ -481,6 +483,90 @@ impl ColumnValueEncoder for ByteArrayEncoder { Ok(()) } + fn count_values_within_byte_budget_gather( + values: &Self::Values, + indices: &[usize], + byte_budget: usize, + ) -> Option { + // `ByteArrayEncoder` only ever writes via `write_gather`, so this + // is the relevant method. + // + // Two-stage walk for the simple offset-buffer byte array types: + // 1. If indices are contiguous, compute the total payload in + // O(1) via a single subtraction on the offsets buffer. + // When the total fits the budget — the overwhelmingly + // common "small values" case — return immediately. + // 2. Otherwise, walk per-value byte sizes from the offsets + // buffer (still cheap, no slice/UTF-8 construction) and + // exit at the first value that pushes the cumulative sum + // past the budget. This bounds skewed distributions: an + // outlier value is caught wherever it lands in the chunk. + let count = match values.data_type() { + DataType::Utf8 => count_within_budget_offsets( + values.as_any().downcast_ref::().unwrap(), + indices, + byte_budget, + ), + DataType::LargeUtf8 => count_within_budget_offsets( + values.as_any().downcast_ref::().unwrap(), + indices, + byte_budget, + ), + DataType::Binary => count_within_budget_offsets( + values.as_any().downcast_ref::().unwrap(), + indices, + byte_budget, + ), + DataType::LargeBinary => count_within_budget_offsets( + values.as_any().downcast_ref::().unwrap(), + indices, + byte_budget, + ), + // View arrays carry each value's length in the low 32 bits of + // its u128 view word, so lengths are scannable without touching + // any data buffer — and the common small-value case skips even + // that scan via an O(1) conservative bound. + DataType::Utf8View => { + let array = values.as_any().downcast_ref::().unwrap(); + count_within_budget_views( + array.views(), + indices, + byte_budget, + max_view_value_len(array.data_buffers()), + ) + } + DataType::BinaryView => { + let array = values.as_any().downcast_ref::().unwrap(); + count_within_budget_views( + array.views(), + indices, + byte_budget, + max_view_value_len(array.data_buffers()), + ) + } + // For arrow Dictionary input, treat every chunk as fitting + // and stay on the batched path. The arrow array being + // Dictionary-encoded in the first place implies its values + // are small enough that dedup is worthwhile, which is the + // opposite of the "5 MiB blob per row" case this fix + // targets. Doing a per-value walk through dict keys (each + // value lookup is keys[i] → values[key] → slice) on every + // chunk costs ~+30-80% vs `main` after writer-dict spill, + // and there is essentially nothing to bound. + DataType::Dictionary(_, _) => indices.len(), + // FixedSizeBinary falls through to the per-value walk via + // `ArrayAccessor::value`. + _ => downcast_op!( + values.data_type(), + values, + count_within_budget_accessor, + indices, + byte_budget + ), + }; + Some(count) + } + fn num_values(&self) -> usize { match &self.dict_encoder { Some(encoder) => encoder.indices.len(), @@ -593,6 +679,141 @@ where } } +/// Cumulative-scan fallback used for byte array types that don't expose +/// a single contiguous offsets buffer — view arrays, dictionary arrays, +/// fixed-size binary. Returns the largest `k` such that the first `k` +/// values picked out by `indices` encode to at most `byte_budget` bytes +/// (or `indices.len()` if they all fit, or `1` if a single value alone +/// exceeds the budget). +/// +/// Free function so it can be used with `downcast_op!`. +fn count_within_budget_accessor(values: T, indices: &[usize], byte_budget: usize) -> usize +where + T: ArrayAccessor + Copy, + T::Item: AsRef<[u8]>, +{ + let mut cum: usize = 0; + for (i, idx) in indices.iter().enumerate() { + let value_len = values.value(*idx).as_ref().len() + std::mem::size_of::(); + cum = cum.saturating_add(value_len); + if cum > byte_budget { + return i.max(1); + } + } + indices.len() +} + +/// Upper bound on any single value's byte length in a view array. +/// +/// An inline view stores at most 12 bytes; an +/// out-of-line view's data is a contiguous slice of exactly one data +/// buffer, so it cannot be longer than the largest data buffer. This is a +/// loose bound (a value is usually far smaller than a whole buffer) but it +/// is O(number of buffers) and always sound. +fn max_view_value_len(buffers: &[Buffer]) -> usize { + /// Bytes that fit inline in a u128 view word (the rest is len + prefix). + const MAX_INLINE_VIEW_LEN: usize = 12; + buffers + .iter() + .map(|b| b.len()) + .max() + .unwrap_or(0) + .max(MAX_INLINE_VIEW_LEN) +} + +/// Two-stage budget count for view arrays (`Utf8View`, `BinaryView`). +/// +/// 1. View arrays have no prefix-sum offsets buffer, so the exact O(1) +/// span subtraction `count_within_budget_offsets` uses is unavailable. +/// But a *conservative* O(1) bound is: every value is at most +/// `max_value_len` bytes, so the whole chunk fits the budget when +/// `n * (max_value_len + 4) <= byte_budget`. This skips the per-value +/// walk for the common small-value case — what view arrays are built +/// for, and exactly the case where there is nothing to bound. +/// 2. Otherwise scan per-value lengths from the low 32 bits of each u128 +/// view word (no data-buffer dereference) and stop at the first value +/// that pushes the cumulative sum past the budget. +fn count_within_budget_views( + views: &[u128], + indices: &[usize], + byte_budget: usize, + max_value_len: usize, +) -> usize { + // Stage 1: O(1) conservative upper bound. + let per_value = max_value_len + std::mem::size_of::(); + if indices.len().saturating_mul(per_value) <= byte_budget { + return indices.len(); + } + // Stage 2: exact per-value scan. + let mut cum: usize = 0; + for (i, idx) in indices.iter().enumerate() { + let len = (views[*idx] as u32) as usize; + cum = cum.saturating_add(len + std::mem::size_of::()); + if cum > byte_budget { + return i.max(1); + } + } + indices.len() +} + +/// Two-stage fast path for `GenericByteArray` +/// (Utf8/LargeUtf8/Binary/LargeBinary). +/// +/// `indices` are assumed sorted ascending — they always are here, since +/// they come from `non_null_indices`, which is built in array order. +/// +/// 1. The span `offsets[last+1] - offsets[first]` is an O(1) upper +/// bound on the chunk's payload: it covers every array position in +/// `[first, last]`, a superset of `indices`. For a non-null chunk +/// `indices` *is* that whole range; for a chunk drawn from a +/// nullable column the skipped positions are nulls, whose offset +/// delta is zero, so the span still equals the exact payload. +/// Either way, if the upper bound fits the budget every value +/// fits — return `indices.len()` with no per-value work. This +/// covers the overwhelmingly common "small values" case for both +/// non-null *and* nullable columns. +/// 2. Otherwise the chunk is genuinely near the budget: walk per-index +/// lengths from the offsets buffer directly (no slice/UTF-8 +/// construction) and stop at the first value that pushes the +/// cumulative sum past the budget. +fn count_within_budget_offsets( + values: &GenericByteArray, + indices: &[usize], + byte_budget: usize, +) -> usize { + if indices.is_empty() { + return 0; + } + let n = indices.len(); + let first = indices[0]; + let last = indices[n - 1]; + let offsets = values.value_offsets(); + let prefix_overhead = std::mem::size_of::(); + + // Stage 1: O(1) span upper bound. Skips Stage 2 in the common case — + // including nullable columns, whose `indices` are sparse. The earlier + // `last - first + 1 == n` contiguity gate forced every nullable chunk + // onto the O(n) Stage 2 walk even though the span check is valid for + // any sorted index set. + if last >= first { + let payload = (offsets[last + 1] - offsets[first]).as_usize(); + if payload + n * prefix_overhead <= byte_budget { + return n; + } + } + + // Stage 2: scan per-index lengths from the offsets buffer. + let mut cum: usize = 0; + for (i, idx) in indices.iter().enumerate() { + let len = (offsets[idx + 1] - offsets[*idx]).as_usize() + prefix_overhead; + cum = cum.saturating_add(len); + if cum > byte_budget { + return i.max(1); + } + } + n +} + /// Computes the min and max for the provided array and indices /// /// This is a free function so it can be used with `downcast_op!` From c2a7fbadd9091925f0917a49fab0f22454f57239 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Mon, 25 May 2026 12:45:47 -0500 Subject: [PATCH 04/15] fix(parquet): bound dictionary page size during dictionary encoding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../src/column/writer/byte_budget_chunker.rs | 67 +++++++++++++------ parquet/tests/arrow_writer_layout.rs | 17 +++-- 2 files changed, 56 insertions(+), 28 deletions(-) diff --git a/parquet/src/column/writer/byte_budget_chunker.rs b/parquet/src/column/writer/byte_budget_chunker.rs index 465568b4caf6..d7276381e655 100644 --- a/parquet/src/column/writer/byte_budget_chunker.rs +++ b/parquet/src/column/writer/byte_budget_chunker.rs @@ -51,6 +51,12 @@ pub(crate) struct ByteBudgetChunker { /// decision short-circuit with no work for every numeric, bool, or /// narrow `FIXED_LEN_BYTE_ARRAY` column. static_always_fits: bool, + /// Configured dictionary page byte limit for the column. + dict_page_byte_limit: usize, + /// As [`Self::static_always_fits`] but for the dictionary page: `true` + /// when one `base_batch_size` mini-batch of this fixed-width type cannot + /// overshoot `dict_page_byte_limit` by more than one mini-batch's worth. + static_dict_always_fits: bool, } impl ByteBudgetChunker { @@ -61,6 +67,7 @@ impl ByteBudgetChunker { base_batch_size: usize, ) -> Self { let page_byte_limit = props.column_data_page_size_limit(descr.path()); + let dict_page_byte_limit = props.column_dictionary_page_size_limit(descr.path()); let static_bytes_per_value = match descr.physical_type() { Type::BOOLEAN => Some(1), Type::INT32 | Type::FLOAT => Some(std::mem::size_of::()), @@ -69,27 +76,36 @@ impl ByteBudgetChunker { Type::FIXED_LEN_BYTE_ARRAY => Some(descr.type_length().max(0) as usize), Type::BYTE_ARRAY => None, }; - let static_always_fits = static_bytes_per_value - .map(|b| b.saturating_mul(base_batch_size) <= page_byte_limit) - .unwrap_or(false); + let static_fits = |limit: usize| { + static_bytes_per_value + .map(|b| b.saturating_mul(base_batch_size) <= limit) + .unwrap_or(false) + }; Self { page_byte_limit, max_def_level: descr.max_def_level(), - static_always_fits, + static_always_fits: static_fits(page_byte_limit), + dict_page_byte_limit, + static_dict_always_fits: static_fits(dict_page_byte_limit), } } /// Decide how many levels at the start of a chunk belong in one - /// mini-batch, so the mini-batch cannot overflow the data page that is - /// currently accumulating value bytes. A returned value smaller than - /// `chunk_size` triggers granular sub-batching in + /// 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`. /// + /// While dictionary-encoding, the data page holds only small RLE indices, + /// but the dictionary page accumulates the distinct values themselves — + /// so it is the dictionary page's remaining budget that must bound the + /// 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, when the column is dictionary-encoding (the data page then - /// holds only small RLE indices, so the plain-encoded byte estimate does - /// not apply), or when the column is a fixed-width type whose - /// mini-batches statically cannot overshoot the data page. + /// 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`. @@ -106,27 +122,34 @@ impl ByteBudgetChunker { if chunk_size == 0 { return chunk_size; } - // While dictionary-encoding, the data page holds only small RLE - // indices — the plain-encoded byte estimate would spuriously shrink - // pages — so stay on the batched fast path. - if encoder.has_dictionary() { - return chunk_size; - } - if self.static_always_fits { - return chunk_size; - } + let budget = if encoder.has_dictionary() { + if self.static_dict_always_fits { + return chunk_size; + } + // 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, + } + } else { + if self.static_always_fits { + return chunk_size; + } + self.page_byte_limit + }; self.byte_budget_sub_batch_size::( values, value_indices, chunk_def, values_offset, chunk_size, - self.page_byte_limit, + budget, ) } /// Inspect value sizes to decide how many of the chunk's values fit in - /// `budget` bytes (the data page remaining budget). + /// `budget` bytes (the data page or dictionary page remaining budget). /// /// `#[inline(never)]` keeps this slow path out of the hot /// `write_batch_internal` loop; numeric and bool columns never reach it. diff --git a/parquet/tests/arrow_writer_layout.rs b/parquet/tests/arrow_writer_layout.rs index b9d997beb289..183590f9b847 100644 --- a/parquet/tests/arrow_writer_layout.rs +++ b/parquet/tests/arrow_writer_layout.rs @@ -408,16 +408,16 @@ fn test_string() { columns: vec![ColumnChunk { pages: vec![ Page { - rows: 130, + rows: 125, page_header_size: 38, - compressed_size: 138, + compressed_size: 114, encoding: Encoding::RLE_DICTIONARY, page_type: PageType::DATA_PAGE, }, Page { - rows: 1250, + rows: 1255, page_header_size: 40, - compressed_size: 10000, + compressed_size: 10040, encoding: Encoding::PLAIN, page_type: PageType::DATA_PAGE, }, @@ -429,10 +429,15 @@ fn test_string() { page_type: PageType::DATA_PAGE, }, ], + // The byte-budget chunker now sub-batches the + // dictionary-encoding phase, so the dictionary page is + // bounded at the 1000-byte limit instead of overshooting + // to 1040 — the dictionary spills one mini-batch earlier + // (125 rows rather than 130). dictionary_page: Some(Page { - rows: 130, + rows: 125, page_header_size: 38, - compressed_size: 1040, + compressed_size: 1000, encoding: Encoding::PLAIN, page_type: PageType::DICTIONARY_PAGE, }), From e05ffca7eda112eb4cfc4961cfa393f4ce9e2a9e Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 28 May 2026 11:29:10 -0500 Subject: [PATCH 05/15] fix(parquet): include boundary value in mini-batch byte-budget count 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) --- parquet/src/arrow/arrow_writer/byte_array.rs | 15 ++++++++------- parquet/src/column/writer/encoder.rs | 9 ++++++--- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index c7e172f1a815..1b36a624588e 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -681,10 +681,11 @@ where /// Cumulative-scan fallback used for byte array types that don't expose /// a single contiguous offsets buffer — view arrays, dictionary arrays, -/// fixed-size binary. Returns the largest `k` such that the first `k` -/// values picked out by `indices` encode to at most `byte_budget` bytes -/// (or `indices.len()` if they all fit, or `1` if a single value alone -/// exceeds the budget). +/// fixed-size binary. Returns `indices.len()` if every value fits the +/// budget, otherwise the smallest `k ≥ 1` whose first `k` values' encoded +/// size first exceeds `byte_budget` — i.e. the boundary value is included +/// so the caller's page-flush check trips immediately on this mini-batch, +/// without leaving a sliver to glue onto the next page. /// /// Free function so it can be used with `downcast_op!`. fn count_within_budget_accessor(values: T, indices: &[usize], byte_budget: usize) -> usize @@ -697,7 +698,7 @@ where let value_len = values.value(*idx).as_ref().len() + std::mem::size_of::(); cum = cum.saturating_add(value_len); if cum > byte_budget { - return i.max(1); + return i + 1; } } indices.len() @@ -750,7 +751,7 @@ fn count_within_budget_views( let len = (views[*idx] as u32) as usize; cum = cum.saturating_add(len + std::mem::size_of::()); if cum > byte_budget { - return i.max(1); + return i + 1; } } indices.len() @@ -808,7 +809,7 @@ fn count_within_budget_offsets( let len = (offsets[idx + 1] - offsets[*idx]).as_usize() + prefix_overhead; cum = cum.saturating_add(len); if cum > byte_budget { - return i.max(1); + return i + 1; } } n diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index 83c08e7423fa..33011bd22f5e 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -307,12 +307,15 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { // pushes us past the budget. This both bounds skewed // distributions (one outlier among small values is caught when // it lands, regardless of position) and short-circuits when an - // early value alone exceeds the budget. + // early value alone exceeds the budget. The boundary value is + // included in the returned count so the caller's page-flush check + // trips on this mini-batch rather than leaving a sliver to glue + // onto the next page (see #9972 discussion). let mut cum: usize = 0; for (i, v) in values[start..end].iter().enumerate() { cum = cum.saturating_add(plain_encoded_byte_size::(v)); if cum > byte_budget { - return Some(i.max(1)); + return Some(i + 1); } } Some(n) @@ -334,7 +337,7 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { let Some(v) = values.get(*idx) else { continue }; cum = cum.saturating_add(plain_encoded_byte_size::(v)); if cum > byte_budget { - return Some(i.max(1)); + return Some(i + 1); } } Some(indices.len()) From 47695880709bfb1ab6e9e698f764574cb4e863ae Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 28 May 2026 11:49:01 -0500 Subject: [PATCH 06/15] test(parquet): update arrow_writer_layout dict-spill expectations 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) --- parquet/tests/arrow_writer_layout.rs | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/parquet/tests/arrow_writer_layout.rs b/parquet/tests/arrow_writer_layout.rs index 183590f9b847..fb4d535a6c98 100644 --- a/parquet/tests/arrow_writer_layout.rs +++ b/parquet/tests/arrow_writer_layout.rs @@ -408,16 +408,16 @@ fn test_string() { columns: vec![ColumnChunk { pages: vec![ Page { - rows: 125, + rows: 126, page_header_size: 38, compressed_size: 114, encoding: Encoding::RLE_DICTIONARY, page_type: PageType::DATA_PAGE, }, Page { - rows: 1255, + rows: 1254, page_header_size: 40, - compressed_size: 10040, + compressed_size: 10032, encoding: Encoding::PLAIN, page_type: PageType::DATA_PAGE, }, @@ -429,15 +429,16 @@ fn test_string() { page_type: PageType::DATA_PAGE, }, ], - // The byte-budget chunker now sub-batches the - // dictionary-encoding phase, so the dictionary page is - // bounded at the 1000-byte limit instead of overshooting - // to 1040 — the dictionary spills one mini-batch earlier - // (125 rows rather than 130). + // The byte-budget chunker sub-batches the dictionary + // phase. The mini-batch deliberately includes the value + // that crosses the 1000-byte limit so the spill triggers + // on this chunk rather than carrying a sliver into the + // next page (see #9972 discussion), giving a 126-row + // dictionary page at 1008 bytes. dictionary_page: Some(Page { - rows: 125, + rows: 126, page_header_size: 38, - compressed_size: 1000, + compressed_size: 1008, encoding: Encoding::PLAIN, page_type: PageType::DICTIONARY_PAGE, }), From c018f5f83355f8ba625cef82bfe066be55a016fb Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 3 Jun 2026 16:20:21 -0400 Subject: [PATCH 07/15] test(parquet): add arrow_writer_layout coverage for dictionary + FixedSizeBinary page sizing Adds two `LayoutTest` cases to `arrow_writer_layout.rs` that exercise byte-budget page-sizing paths introduced in #9972 through the real `ArrowWriter` user path: - `test_dictionary`: an arrow `DictionaryArray` 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`. 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) --- parquet/tests/arrow_writer_layout.rs | 117 +++++++++++++++++++++++++++ 1 file changed, 117 insertions(+) diff --git a/parquet/tests/arrow_writer_layout.rs b/parquet/tests/arrow_writer_layout.rs index fb4d535a6c98..cd1225f64f08 100644 --- a/parquet/tests/arrow_writer_layout.rs +++ b/parquet/tests/arrow_writer_layout.rs @@ -20,6 +20,8 @@ use arrow::array::{Int32Array, StringArray}; use arrow::record_batch::RecordBatch; use arrow_array::builder::{Int32Builder, ListBuilder}; +use arrow_array::types::Int32Type; +use arrow_array::{DictionaryArray, FixedSizeBinaryArray}; use bytes::Bytes; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::{ArrowReaderOptions, ParquetRecordBatchReaderBuilder}; @@ -605,3 +607,118 @@ fn test_per_column_data_page_size_limit() { assert_eq!(col_a_page_count, 16); assert_eq!(col_b_page_count, 1); } + +#[test] +fn test_fixed_size_binary() { + // FixedSizeBinary values larger than the data page byte limit. + // + // The arrow writer routes a (non-dictionary) FixedSizeBinary column + // through the *generic* column writer + // (`ColumnValueEncoderImpl`), not the byte-array + // encoder. Its byte-budget chunker uses the plain-encoded value size — + // `type_length` bytes, with no length prefix for FIXED_LEN_BYTE_ARRAY — + // to cap each data page. With 1 KiB values and a 4 KiB page limit, a + // mini-batch grows until its cumulative bytes first cross the limit + // (the boundary value is included), so 5 values land per page (5120 B) + // instead of buffering all 64 values into one ~64 KiB page. + let value_size = 1024usize; + let num_rows = 64usize; + let values: Vec = (0..num_rows) + .flat_map(|i| vec![i as u8; value_size]) + .collect(); + let array = + Arc::new(FixedSizeBinaryArray::try_new(value_size as i32, values.into(), None).unwrap()) as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(4096) + .set_write_page_header_statistics(true) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + // 12 pages of 5 values (5 * 1024 = 5120 B, the boundary + // value pushes each page just past the 4096 B limit) plus + // a final page with the remaining 4 values. + pages: (0..12) + .map(|_| Page { + rows: 5, + page_header_size: 157, + compressed_size: 5120, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }) + .chain(std::iter::once(Page { + rows: 4, + page_header_size: 157, + compressed_size: 4096, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + })) + .collect(), + dictionary_page: None, + }], + }], + }, + }); +} + +#[test] +fn test_dictionary() { + // Arrow `DictionaryArray` input. + // + // The byte-budget chunker's gather path has a dedicated arm for + // dictionary-typed arrow input that reports "everything fits" + // (`indices.len()`), keeping the column on the batched fast path: an + // arrow array that is *already* dictionary-encoded implies its values + // are small enough that dedup is worthwhile — the opposite of the large + // blob case the page-size fix targets — so there is nothing to bound and + // walking dict keys per chunk would only cost time. This test drives that + // arm and confirms the column still lays out as a normal dictionary + // column (one small dictionary page + one RLE_DICTIONARY data page). + let num_rows = 2000; + let dict_values = StringArray::from_iter_values(["alpha", "beta", "gamma", "delta"]); + let keys = Int32Array::from_iter_values((0..num_rows as i32).map(|i| i % 4)); + let array = Arc::new(DictionaryArray::::try_new(keys, Arc::new(dict_values)).unwrap()) + as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(true) + .set_dictionary_page_size_limit(1000) + .set_data_page_size_limit(1000) + .set_write_batch_size(10) + .set_write_page_header_statistics(true) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + pages: vec![Page { + rows: 2000, + page_header_size: 40, + compressed_size: 505, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }], + dictionary_page: Some(Page { + rows: 4, + page_header_size: 38, + compressed_size: 35, + encoding: Encoding::PLAIN, + page_type: PageType::DICTIONARY_PAGE, + }), + }], + }], + }, + }); +} + From 85741877b187fc9a9f82f703d56e387b91eeab1e Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Mon, 25 May 2026 12:44:25 -0500 Subject: [PATCH 08/15] test(parquet): add page-size regression tests for large values 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) --- parquet/src/arrow/arrow_writer/mod.rs | 239 +++++++++++++ parquet/src/column/writer/mod.rs | 494 ++++++++++++++++++++++++++ 2 files changed, 733 insertions(+) diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 79542caed9b7..7c3152d05266 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -4904,6 +4904,245 @@ mod tests { assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4); } + #[test] + fn test_arrow_writer_caps_page_size_for_large_strings() { + // End-to-end coverage for `ByteArrayEncoder::estimated_value_bytes_gather` + // and the offsets-buffer fast path in `estimate_byte_size_offsets`. + // + // The standard column-writer regression test covers the same + // logic for the non-arrow path; this one drives writes through + // `ArrowWriter`, which is what real users hit and which uses + // `write_gather` with `non_null_indices` (always contiguous for + // a non-null column). + let value_size = 64 * 1024; + let page_byte_limit = 16 * 1024; + let num_rows = 32; + + let schema = Arc::new(Schema::new(vec![Field::new( + "col", + ArrowDataType::Utf8, + false, + )])); + let strings: Vec = (0..num_rows).map(|_| "x".repeat(value_size)).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StringArray::from(strings)) as _], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(page_byte_limit) + .build(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + let data = Bytes::from(writer.into_inner().unwrap()); + + let mut metadata = ParquetMetaDataReader::new(); + metadata.try_parse(&data).unwrap(); + let metadata = metadata.finish().unwrap(); + let col_meta = metadata.row_group(0).column(0); + let mut reader = + SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap(); + + let mut data_pages = Vec::new(); + while let Some(page) = reader.get_next_page().unwrap() { + if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) { + data_pages.push(page.buffer().len()); + } + } + + assert!( + data_pages.len() >= num_rows / 2, + "expected at least {} data pages for {num_rows} large strings via ArrowWriter, got {} ({data_pages:?})", + num_rows / 2, + data_pages.len(), + ); + for size in &data_pages { + assert!( + *size <= value_size + 64, + "page size {size} exceeds one-value bound; pages {data_pages:?}" + ); + } + } + + #[test] + fn test_arrow_writer_granular_mode_roundtrip() { + // Granular mode subdivides chunks and writes more pages than + // `main`. Make sure the data we write back is bit-identical to + // what went in — page-count assertions elsewhere only prove + // pages were cut, not that the encoded data is correct. + // + // Mix value sizes so that the cumulative-byte-budget cutoff + // lands mid-chunk, exercising both batched and granular paths + // within the same `write_batch_internal` call. + let small = "tiny".to_string(); + let big = "x".repeat(64 * 1024); + let strings: Vec = (0..256) + .map(|i| { + if i % 16 == 0 { + big.clone() + } else { + small.clone() + } + }) + .collect(); + + let schema = Arc::new(Schema::new(vec![Field::new( + "col", + ArrowDataType::Utf8, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StringArray::from(strings.clone())) as _], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(16 * 1024) + .build(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + let data = Bytes::from(writer.into_inner().unwrap()); + + let mut reader = ParquetRecordBatchReader::try_new(data, 1024).unwrap(); + let read = reader.next().unwrap().unwrap(); + assert!(reader.next().is_none(), "expected one batch"); + let col = read + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(col.len(), strings.len()); + for (i, expected) in strings.iter().enumerate() { + assert_eq!( + col.value(i), + expected.as_str(), + "value mismatch at index {i}" + ); + } + } + + #[test] + fn test_arrow_writer_caps_page_size_for_large_string_view() { + // Coverage for the view/dict fallback in + // `ByteArrayEncoder::count_values_within_byte_budget_gather`. + // `Utf8View` arrays don't expose a single contiguous offsets + // buffer, so the gather method falls back to the per-value walk + // via `ArrayAccessor::value`. + use arrow_array::StringViewArray; + let value_size = 64 * 1024; + let page_byte_limit = 16 * 1024; + let num_rows = 32usize; + + let schema = Arc::new(Schema::new(vec![Field::new( + "col", + ArrowDataType::Utf8View, + false, + )])); + let strings: Vec = (0..num_rows).map(|_| "x".repeat(value_size)).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StringViewArray::from_iter_values( + strings.iter().map(|s| s.as_str()), + )) as _], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(page_byte_limit) + .build(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + let data = Bytes::from(writer.into_inner().unwrap()); + + let mut metadata = ParquetMetaDataReader::new(); + metadata.try_parse(&data).unwrap(); + let metadata = metadata.finish().unwrap(); + let col_meta = metadata.row_group(0).column(0); + let mut reader = + SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap(); + + let mut data_pages = Vec::new(); + while let Some(page) = reader.get_next_page().unwrap() { + if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) { + data_pages.push(page.buffer().len()); + } + } + assert!( + data_pages.len() >= num_rows / 2, + "expected pages bounded by byte budget for Utf8View, got {data_pages:?}" + ); + for size in &data_pages { + assert!( + *size <= value_size + 64, + "page size {size} exceeds one-value bound; pages {data_pages:?}" + ); + } + } + + #[test] + fn test_arrow_writer_all_null_string_column() { + // The `LevelDataRef::value_count` Uniform branch with + // `value != max_def` (entirely-null chunk) must return 0 so the + // sub-batch sizer short-circuits to batch mode without trying + // to estimate byte budgets for non-existent values. + let num_rows = 1024; + let schema = Arc::new(Schema::new(vec![Field::new( + "col", + ArrowDataType::Utf8, + true, + )])); + let nulls: Vec> = vec![None; num_rows]; + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StringArray::from(nulls)) as _], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(16 * 1024) + .build(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + let data = Bytes::from(writer.into_inner().unwrap()); + + // Re-parse the file: row group has one column, every row is + // null, all data pages report `num_rows / page_count` rows. + let mut metadata = ParquetMetaDataReader::new(); + metadata.try_parse(&data).unwrap(); + let metadata = metadata.finish().unwrap(); + let row_group = metadata.row_group(0); + let col_meta = row_group.column(0); + assert_eq!(row_group.num_rows() as usize, num_rows); + // Statistics record `null_count = num_rows` — proves every value + // was written as null. + if let Some(stats) = col_meta.statistics() { + assert_eq!( + stats.null_count_opt().unwrap_or(0) as usize, + num_rows, + "expected all-null column to report null_count = num_rows" + ); + } + + let mut reader = + SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap(); + let mut total_values = 0u32; + while let Some(page) = reader.get_next_page().unwrap() { + if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) { + total_values += page.num_values(); + } + } + assert_eq!( + total_values as usize, num_rows, + "expected every level position to be represented in some page" + ); + } + struct WriteBatchesShape { num_batches: usize, rows_per_batch: usize, diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index 4e53230bbf89..97ccf13a316b 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -2676,6 +2676,500 @@ mod tests { assert_eq!(other_values, vec![10]); } + #[test] + fn test_column_writer_caps_page_size_for_large_byte_array_values() { + // Regression: the post-write data page byte limit check only fires + // at mini-batch boundaries, so a 1024-row mini-batch of multi-MiB + // BYTE_ARRAY values used to buffer multiple GiB into a single page + // before the limit was even consulted. With the threshold-based + // granular mode this batch should split into ~one page per value. + let value_size = 64 * 1024; // 64 KiB per value + let page_byte_limit = 16 * 1024; // 16 KiB page limit + let num_rows = 64; + + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let props = Arc::new( + WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_data_page_size_limit(page_byte_limit) + // Default write_batch_size (1024) — without the fix this + // buffers the entire input into a single ~4 MiB page. + .build(), + ); + + let mut data = Vec::with_capacity(num_rows); + for i in 0..num_rows { + data.push(ByteArray::from(vec![i as u8; value_size])); + } + + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + writer.write_batch(&data, None, None).unwrap(); + let r = writer.close().unwrap(); + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let mut page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(), + ); + + let mut data_page_sizes = Vec::new(); + let mut data_page_value_counts = Vec::new(); + while let Some(page) = page_reader.get_next_page().unwrap() { + if matches!( + page.page_type(), + PageType::DATA_PAGE | PageType::DATA_PAGE_V2 + ) { + data_page_sizes.push(page.buffer().len()); + data_page_value_counts.push(page.num_values()); + } + } + + // Every value must end up somewhere. + assert_eq!( + data_page_value_counts.iter().sum::() as usize, + num_rows + ); + // Without the fix this assertion fired with one ~4 MiB page; the + // threshold splits the input so that no page holds more than a + // single oversized value's worth of bytes. + assert!( + data_page_sizes.len() >= num_rows / 2, + "expected pages to be cut close to one per value, got {} pages for sizes {:?}", + data_page_sizes.len(), + data_page_sizes, + ); + // Each page must be bounded by roughly one value's worth of bytes; + // parquet allows a single oversized value to occupy a page by + // itself but never lets us pile many of them together. + for size in &data_page_sizes { + assert!( + *size <= value_size + 64, + "page size {size} exceeds one-value bound ({}B) — pages {:?}", + value_size + 64, + data_page_sizes, + ); + } + } + + #[test] + fn test_column_writer_caps_page_size_for_large_values_in_list() { + // Coverage for the Materialized-rep branch of + // `write_granular_chunk`. The flat-column regression test + // exercises the per-level step; this exercises the + // record-by-record step used when rep levels are present. + // + // Column is `list` (max_def = 1, max_rep = 1) + // with 3 records of 3 large blobs each. The page byte limit is + // smaller than a single blob, so granular mode kicks in, and the + // Materialized-rep arm of `write_granular_chunk` steps from one + // `rep == 0` boundary to the next so a record never spans pages. + let value_size = 32 * 1024; + let page_byte_limit = 16 * 1024; + let values_per_record = 3; + let num_records = 3; + let num_values = values_per_record * num_records; + + // rep levels: 0, 1, 1, 0, 1, 1, 0, 1, 1 + let mut rep_levels = Vec::with_capacity(num_values); + for _ in 0..num_records { + rep_levels.push(0i16); + rep_levels.extend(std::iter::repeat_n(1i16, values_per_record - 1)); + } + let def_levels = vec![1i16; num_values]; + + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let props = Arc::new( + WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_data_page_size_limit(page_byte_limit) + .build(), + ); + + let mut data = Vec::with_capacity(num_values); + for i in 0..num_values { + data.push(ByteArray::from(vec![i as u8; value_size])); + } + + let mut writer = get_test_column_writer::(page_writer, 1, 1, props); + writer + .write_batch(&data, Some(&def_levels), Some(&rep_levels)) + .unwrap(); + let r = writer.close().unwrap(); + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let mut page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(), + ); + + let mut data_pages = Vec::new(); + while let Some(page) = page_reader.get_next_page().unwrap() { + if matches!( + page.page_type(), + PageType::DATA_PAGE | PageType::DATA_PAGE_V2 + ) { + data_pages.push((page.buffer().len(), page.num_values())); + } + } + + // The Materialized-rep arm groups levels by record, and each + // record's bytes blow the page byte limit on its own, so we get + // exactly one page per record. + assert_eq!( + data_pages.len(), + num_records, + "expected one data page per record, got {data_pages:?}" + ); + for (bytes, n_values) in &data_pages { + assert_eq!( + *n_values as usize, values_per_record, + "each page must hold a whole record's leaves, got {data_pages:?}" + ); + // Each page is one full record (its leaves cannot be split), + // so allow up to `values_per_record` blobs of payload plus a + // small fudge for level encoding overhead. + let upper_bound = values_per_record * (value_size + 16); + assert!( + *bytes <= upper_bound, + "page size {bytes} exceeds whole-record bound ({upper_bound}); pages {data_pages:?}" + ); + } + } + + #[test] + fn test_column_writer_caps_page_size_with_nullable_large_values() { + // Coverage for `LevelDataRef::value_count` on Materialized def + // levels: a nullable column with mixed nulls and large values. + // `value_count` must return the actual non-null count so the + // byte estimate reflects bytes that will actually be written, + // not the level count. + let value_size = 32 * 1024; + let page_byte_limit = 16 * 1024; + let num_levels = 32; + + // Alternating null / non-null: 16 nulls and 16 values. + let def_levels: Vec = (0..num_levels as i16).map(|i| i % 2).collect(); + let num_values = def_levels.iter().filter(|&&d| d == 1).count(); + + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let props = Arc::new( + WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_data_page_size_limit(page_byte_limit) + .build(), + ); + + let mut data = Vec::with_capacity(num_values); + for i in 0..num_values { + data.push(ByteArray::from(vec![i as u8; value_size])); + } + + let mut writer = get_test_column_writer::(page_writer, 1, 0, props); + writer.write_batch(&data, Some(&def_levels), None).unwrap(); + let r = writer.close().unwrap(); + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let mut page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(), + ); + + let mut data_pages = Vec::new(); + while let Some(page) = page_reader.get_next_page().unwrap() { + if matches!( + page.page_type(), + PageType::DATA_PAGE | PageType::DATA_PAGE_V2 + ) { + data_pages.push(page.buffer().len()); + } + } + + // With 16 actual values of 32 KiB each and a 16 KiB page limit, + // every non-null value should get its own page (plus possibly + // adjacent nulls). At minimum, the number of pages must be + // roughly the value count, not 1 (which is what `main` produced). + assert!( + data_pages.len() >= num_values / 2, + "expected at least {} pages for {num_values} large values, got {} pages: {data_pages:?}", + num_values / 2, + data_pages.len(), + ); + // No page contains more than ~one value's worth of payload bytes. + for size in &data_pages { + assert!( + *size <= value_size + 64, + "page size {size} exceeds one-value bound; pages {data_pages:?}" + ); + } + } + + #[test] + fn test_column_writer_dict_enabled_large_values_post_spill() { + // While dictionary encoding is active, `has_dictionary()` short- + // circuits `estimated_value_bytes` — the byte estimate is plain- + // encoded size but dict-encoded pages only store small RLE + // indices, so we'd otherwise shrink pages spuriously. Once the + // dictionary spills (each value is large + unique), plain + // encoding takes over and the byte-budget sub-batch kicks in. + // + // This test makes sure the writer survives that transition and + // produces bounded pages thereafter. + let value_size = 64 * 1024; + let page_byte_limit = 16 * 1024; + let num_rows = 32; + + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let props = Arc::new( + WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(true) + // Force a small dict so it spills quickly even though + // each value here is unique. + .set_dictionary_page_size_limit(1024) + .set_data_page_size_limit(page_byte_limit) + // Small mini-batches so dict fallback happens part-way + // through the input, leaving subsequent mini-batches to + // exercise the post-spill plain-encoding path that the + // page-size fix actually targets. + .set_write_batch_size(4) + .build(), + ); + + let mut data = Vec::with_capacity(num_rows); + for i in 0..num_rows { + data.push(ByteArray::from(vec![i as u8; value_size])); + } + + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + writer.write_batch(&data, None, None).unwrap(); + let r = writer.close().unwrap(); + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let mut page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(), + ); + + let mut data_pages = Vec::new(); + while let Some(page) = page_reader.get_next_page().unwrap() { + if matches!( + page.page_type(), + PageType::DATA_PAGE | PageType::DATA_PAGE_V2 + ) { + data_pages.push(page.buffer().len()); + } + } + + // After spill, plain encoding writes one ~64 KiB value per page. + // Without the fix, post-spill writes still buffered all 32 + // values into a single ~2 MiB page. + assert!( + data_pages.len() >= num_rows / 2, + "expected >= {} data pages after dict spill, got {} ({data_pages:?})", + num_rows / 2, + data_pages.len(), + ); + for size in &data_pages { + assert!( + *size <= value_size + 64, + "page size {size} exceeds one-value bound; pages {data_pages:?}" + ); + } + } + + #[test] + fn test_column_writer_caps_dictionary_page_size() { + // A column of large *distinct* values with dictionary encoding on: + // the dictionary page accumulates the values themselves, and its + // spill check runs only once per mini-batch. Without bounding the + // dictionary-encoding mini-batch, one `write_batch_size` mini-batch + // would intern `write_batch_size * value_size` bytes into the + // dictionary page before the check fires (~16 MiB here). The chunker + // must sub-batch the dictionary-encoding phase too. + let value_size = 8 * 1024; + let dict_page_limit = 64 * 1024; + let num_rows = 2048; + + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let props = Arc::new( + WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(true) + .set_dictionary_page_size_limit(dict_page_limit) + .build(), + ); + + let mut data = Vec::with_capacity(num_rows); + for i in 0..num_rows { + // each value distinct, so the dictionary cannot dedup them + let mut v = vec![0u8; value_size]; + v[..8].copy_from_slice(&(i as u64).to_le_bytes()); + data.push(ByteArray::from(v)); + } + + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + writer.write_batch(&data, None, None).unwrap(); + let r = writer.close().unwrap(); + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let mut page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(), + ); + + let mut dict_page_size = 0; + while let Some(page) = page_reader.get_next_page().unwrap() { + if page.page_type() == PageType::DICTIONARY_PAGE { + dict_page_size = dict_page_size.max(page.buffer().len()); + } + } + + assert!( + dict_page_size > 0, + "expected the column to dictionary-encode" + ); + // Bounded near the limit (~2x from the post-mini-batch check). Before + // the fix the dictionary page reached num_rows * value_size (~16 MiB, + // 256x the limit). + assert!( + dict_page_size <= 3 * dict_page_limit, + "dictionary page {dict_page_size} exceeds 3x the {dict_page_limit} limit", + ); + } + + #[test] + fn test_column_writer_caps_page_size_for_fixed_len_byte_array() { + // Coverage for `ParquetValueType::byte_size` override on + // `FixedLenByteArray`. With `type_length = 1`, each plain-encoded + // value is one byte, so a 4-byte page byte limit forces the + // sub-batch sizer to write ~4 values per page rather than one + // page for the whole batch. + let page_byte_limit = 4; + let num_values = 128; + + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let props = Arc::new( + WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_data_page_size_limit(page_byte_limit) + .build(), + ); + + let mut data = Vec::with_capacity(num_values); + for i in 0..num_values { + let mut fla = FixedLenByteArray::default(); + fla.set_data(Bytes::from(vec![i as u8])); + data.push(fla); + } + + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + writer.write_batch(&data, None, None).unwrap(); + let r = writer.close().unwrap(); + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let mut page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(), + ); + + let mut data_pages = Vec::new(); + while let Some(page) = page_reader.get_next_page().unwrap() { + if matches!( + page.page_type(), + PageType::DATA_PAGE | PageType::DATA_PAGE_V2 + ) { + data_pages.push(page.buffer().len()); + } + } + + // Without the fix this is a single 128-byte page; with the fix + // the byte budget caps each page at ~`page_byte_limit` bytes. + assert!( + data_pages.len() >= num_values / 8, + "expected pages capped by byte budget, got {data_pages:?}" + ); + for size in &data_pages { + assert!( + *size <= page_byte_limit * 4, + "page size {size} larger than expected; pages {data_pages:?}" + ); + } + } + #[test] fn test_bool_statistics() { let stats = statistics_roundtrip::(&[true, false, false, true]); From e20609a56949f4b5a6bdce214cc41afcfaaf1411 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Mon, 25 May 2026 12:45:28 -0500 Subject: [PATCH 09/15] fix(parquet): bound data page byte size for large variable-width values MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../src/column/writer/byte_budget_chunker.rs | 183 ++++++++++++++++++ parquet/src/column/writer/encoder.rs | 128 ++++++++++++ parquet/src/column/writer/mod.rs | 167 +++++++++++++++- 3 files changed, 472 insertions(+), 6 deletions(-) create mode 100644 parquet/src/column/writer/byte_budget_chunker.rs diff --git a/parquet/src/column/writer/byte_budget_chunker.rs b/parquet/src/column/writer/byte_budget_chunker.rs new file mode 100644 index 000000000000..465568b4caf6 --- /dev/null +++ b/parquet/src/column/writer/byte_budget_chunker.rs @@ -0,0 +1,183 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Decides how many levels of a chunk to write as one mini-batch so that +//! the resulting data page stays within `data_page_size_limit`. +//! +//! The parquet column writer checks the data page byte limit only *after* +//! each mini-batch finishes writing. Mini-batches are sized in rows +//! (`write_batch_size`, default 1024), so for BYTE_ARRAY columns whose +//! values are large (e.g. multi-MiB blobs) a single mini-batch can buffer +//! GiB into one page before the limit is consulted. +//! +//! This module isolates the per-chunk decision that prevents that: given a +//! chunk's level data and the input values, pick the largest `sub_batch_size` +//! such that one mini-batch fits in one page byte budget. For the +//! overwhelmingly common case (small or fixed-width values) the answer is +//! just `chunk_size` and the decision is O(1) on the column type — only +//! when the input might overflow does the chunker consult the encoder's +//! byte estimate. + +use crate::basic::Type; +use crate::column::writer::LevelDataRef; +use crate::column::writer::encoder::ColumnValueEncoder; +use crate::file::properties::WriterProperties; +use crate::schema::types::ColumnDescriptor; + +/// Picks byte-budget-aware mini-batch sizes for one column. +pub(crate) struct ByteBudgetChunker { + /// Configured data page byte limit for the column. + page_byte_limit: usize, + /// Max definition level of the column; a level equal to this marks a + /// present (non-null) leaf value. Used to count values per chunk. + max_def_level: i16, + /// `true` when no chunk of `base_batch_size` values can ever overflow + /// `page_byte_limit` regardless of input. Set once at column open from + /// the physical type's known per-value byte size; lets the per-chunk + /// decision short-circuit with no work for every numeric, bool, or + /// narrow `FIXED_LEN_BYTE_ARRAY` column. + static_always_fits: bool, +} + +impl ByteBudgetChunker { + #[inline] + pub(crate) fn new( + descr: &ColumnDescriptor, + props: &WriterProperties, + base_batch_size: usize, + ) -> Self { + let page_byte_limit = props.column_data_page_size_limit(descr.path()); + let static_bytes_per_value = match descr.physical_type() { + Type::BOOLEAN => Some(1), + Type::INT32 | Type::FLOAT => Some(std::mem::size_of::()), + Type::INT64 | Type::DOUBLE => Some(std::mem::size_of::()), + Type::INT96 => Some(12), + Type::FIXED_LEN_BYTE_ARRAY => Some(descr.type_length().max(0) as usize), + Type::BYTE_ARRAY => None, + }; + let static_always_fits = static_bytes_per_value + .map(|b| b.saturating_mul(base_batch_size) <= page_byte_limit) + .unwrap_or(false); + Self { + page_byte_limit, + max_def_level: descr.max_def_level(), + static_always_fits, + } + } + + /// Decide how many levels at the start of a chunk belong in one + /// mini-batch, so the mini-batch cannot overflow the data page that is + /// currently accumulating value bytes. A returned value smaller than + /// `chunk_size` triggers granular sub-batching in + /// `write_batch_internal`. + /// + /// Returns `chunk_size` immediately (no value inspection) when the chunk + /// is empty, when the column is dictionary-encoding (the data page then + /// holds only small RLE indices, so the plain-encoded byte estimate does + /// not apply), or when the column is a fixed-width type whose + /// mini-batches statically cannot overshoot the data page. + /// + /// `#[inline]`: this is a tiny per-chunk dispatcher; the actual byte + /// inspection lives in the out-of-line `byte_budget_sub_batch_size`. + #[inline] + pub(crate) fn pick_sub_batch_size( + &self, + encoder: &E, + values: &E::Values, + value_indices: Option<&[usize]>, + chunk_def: LevelDataRef<'_>, + values_offset: usize, + chunk_size: usize, + ) -> usize { + if chunk_size == 0 { + return chunk_size; + } + // While dictionary-encoding, the data page holds only small RLE + // indices — the plain-encoded byte estimate would spuriously shrink + // pages — so stay on the batched fast path. + if encoder.has_dictionary() { + return chunk_size; + } + if self.static_always_fits { + return chunk_size; + } + self.byte_budget_sub_batch_size::( + values, + value_indices, + chunk_def, + values_offset, + chunk_size, + self.page_byte_limit, + ) + } + + /// Inspect value sizes to decide how many of the chunk's values fit in + /// `budget` bytes (the data page remaining budget). + /// + /// `#[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( + &self, + values: &E::Values, + value_indices: Option<&[usize]>, + chunk_def: LevelDataRef<'_>, + values_offset: usize, + chunk_size: usize, + budget: usize, + ) -> usize { + // 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; + } + // Ask the encoder how many of the next values fit in one page byte + // budget. Dispatch on whether the caller supplied gather indices; + // this mirrors how `write_mini_batch` picks `write_gather` vs + // `write`. + let fit = match value_indices { + Some(idx) => { + let end = (values_offset + vals_in_chunk).min(idx.len()); + let start = values_offset.min(end); + E::count_values_within_byte_budget_gather(values, &idx[start..end], budget) + } + None => { + E::count_values_within_byte_budget(values, values_offset, vals_in_chunk, budget) + } + }; + match fit { + None => chunk_size, + 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 + } else { + (values_per_subbatch * chunk_size) + .div_ceil(vals_in_chunk) + .max(1) + }; + chunk_size.min(levels_per_subbatch.max(1)) + } + } + } +} diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index 2ea3376ae708..83c08e7423fa 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -90,6 +90,42 @@ pub trait ColumnValueEncoder { /// Write the values at the indexes in `indices` to this [`ColumnValueEncoder`] fn write_gather(&mut self, values: &Self::Values, indices: &[usize]) -> Result<()>; + /// Returns the largest `k` such that the first `k` values in + /// `values[offset..offset + len]` encode to at most `byte_budget` + /// bytes — i.e. how many values fit in a single page byte budget. + /// + /// Returns `len` if every value fits. Returns at least 1 if a single + /// value alone exceeds the budget, matching parquet's "at least one + /// value per data page" rule. + /// + /// `None` means "no cheap estimate available"; the caller stays on + /// the batched fast path and lets the post-write + /// `should_add_data_page` check handle bounding. + /// + /// Implementations should short-circuit aggressively: the typical + /// case is "everything fits, return `len`", and the next-most-common + /// case is "one wide value, return 1." The variable-width walk only + /// needs to be precise when the chunk is genuinely near the budget. + fn count_values_within_byte_budget( + _values: &Self::Values, + _offset: usize, + _len: usize, + _byte_budget: usize, + ) -> Option { + None + } + + /// As [`Self::count_values_within_byte_budget`] but using gather + /// `indices` rather than a contiguous range. Returns the number of + /// `indices` that fit, not the maximum index value. + fn count_values_within_byte_budget_gather( + _values: &Self::Values, + _indices: &[usize], + _byte_budget: usize, + ) -> Option { + None + } + /// Returns the number of buffered values fn num_values(&self) -> usize; @@ -247,6 +283,63 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { self.write_slice(&slice) } + fn count_values_within_byte_budget( + values: &[T::T], + offset: usize, + len: usize, + byte_budget: usize, + ) -> Option { + // Clamp so that a caller-supplied `len` that overruns the input + // (e.g. a level/value mismatch the encoder will reject later) + // returns an estimate instead of panicking here. + let end = (offset + len).min(values.len()); + let start = offset.min(end); + let n = end - start; + // Fixed-size physical types have a constant per-value byte cost, + // so the answer is one division — no need to walk the slice. + let phys = ::PHYSICAL_TYPE; + if phys != Type::BYTE_ARRAY && phys != Type::FIXED_LEN_BYTE_ARRAY { + let per = std::mem::size_of::().max(1); + let fits = (byte_budget / per).max(1); + return Some(fits.min(n)); + } + // Variable-width: scan, accumulate, exit at the first value that + // pushes us past the budget. This both bounds skewed + // distributions (one outlier among small values is caught when + // it lands, regardless of position) and short-circuits when an + // early value alone exceeds the budget. + let mut cum: usize = 0; + for (i, v) in values[start..end].iter().enumerate() { + cum = cum.saturating_add(plain_encoded_byte_size::(v)); + if cum > byte_budget { + return Some(i.max(1)); + } + } + Some(n) + } + + fn count_values_within_byte_budget_gather( + values: &[T::T], + indices: &[usize], + byte_budget: usize, + ) -> Option { + let phys = ::PHYSICAL_TYPE; + if phys != Type::BYTE_ARRAY && phys != Type::FIXED_LEN_BYTE_ARRAY { + let per = std::mem::size_of::().max(1); + let fits = (byte_budget / per).max(1); + return Some(fits.min(indices.len())); + } + let mut cum: usize = 0; + for (i, idx) in indices.iter().enumerate() { + let Some(v) = values.get(*idx) else { continue }; + cum = cum.saturating_add(plain_encoded_byte_size::(v)); + if cum > byte_budget { + return Some(i.max(1)); + } + } + Some(indices.len()) + } + fn num_values(&self) -> usize { self.num_values } @@ -411,3 +504,38 @@ where } } } + +/// Plain-encoded byte cost of a single value of type `T::T`. +/// +/// Derived from [`ParquetValueType::dict_encoding_size`] so we don't add a +/// parallel per-value-size hook to the trait. The components returned by +/// `dict_encoding_size` are `(per-value overhead, value-bytes)`. For +/// plain encoding the on-disk layout is: +/// +/// - `BYTE_ARRAY`: 4-byte length prefix + payload bytes = `overhead + bytes`. +/// - `FIXED_LEN_BYTE_ARRAY`: raw bytes only, taken from the type descriptor's +/// `type_length`. The value's own `dict_encoding_size` reports the length +/// prefix, which is irrelevant for plain FLBA encoding; the encoder passes +/// `type_length` directly. +/// - Everything else (numeric / bool): a constant per-value size; the caller +/// already short-circuits these via `mem::size_of::()` before +/// touching this function, so this branch is unreachable in practice and +/// we fall back to `overhead` defensively. +/// +/// See `dict_encoder.rs::push` (line ~52) for the matching dispatch. +/// +/// Placed at the end of the module deliberately. Inserting it above the +/// `ColumnValueEncoder` trait shifts the trait and `ColumnValueEncoderImpl` +/// within the compiled module enough to perturb downstream code placement, +/// which measurably regresses unrelated arrow-writer string benchmarks +/// (~5-9% on `string` / `string_and_binary_view`). Defining it last keeps +/// the hot encoder code at the offsets it has on `main`. +#[inline] +fn plain_encoded_byte_size(value: &T::T) -> usize { + let (overhead, bytes) = value.dict_encoding_size(); + match ::PHYSICAL_TYPE { + Type::BYTE_ARRAY => overhead + bytes, + Type::FIXED_LEN_BYTE_ARRAY => bytes, + _ => overhead, + } +} diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index 97ccf13a316b..52e2ea06a50c 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -49,8 +49,11 @@ use crate::file::properties::{ use crate::file::statistics::{Statistics, ValueStatistics}; use crate::schema::types::{ColumnDescPtr, ColumnDescriptor}; +mod byte_budget_chunker; pub(crate) mod encoder; +use byte_budget_chunker::ByteBudgetChunker; + macro_rules! downcast_writer { ($e:expr, $i:ident, $b:expr) => { match $e { @@ -374,6 +377,24 @@ impl<'a> LevelDataRef<'a> { Self::Uniform { value, .. } => Self::Uniform { value, count: len }, } } + + /// Count of positions in this slice that represent an actual value + /// (definition level equal to `max_def`). `Absent` means the column has + /// `max_def == 0` and every position is a value, so the implicit count + /// is the caller-supplied `total`. + pub(crate) fn value_count(self, total: usize, max_def: i16) -> usize { + match self { + Self::Absent => total, + Self::Materialized(values) => values.iter().filter(|&&d| d == max_def).count(), + Self::Uniform { value, count } => { + if value == max_def { + count + } else { + 0 + } + } + } + } } /// Typed column writer for a primitive column. @@ -545,6 +566,7 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { } else { self.props.write_batch_size() }; + let chunker = ByteBudgetChunker::new(&self.descr, &self.props, base_batch_size); while levels_offset < num_levels { let mut end_offset = num_levels.min(levels_offset + base_batch_size); @@ -555,14 +577,39 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { } } - values_offset += self.write_mini_batch( + let chunk_size = end_offset - levels_offset; + let chunk_def = def_levels.slice(levels_offset, chunk_size); + let chunk_rep = rep_levels.slice(levels_offset, chunk_size); + + let sub_batch_size = chunker.pick_sub_batch_size( + &self.encoder, values, - values_offset, value_indices, - end_offset - levels_offset, - def_levels.slice(levels_offset, end_offset - levels_offset), - rep_levels.slice(levels_offset, end_offset - levels_offset), - )?; + chunk_def, + values_offset, + 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, + )?; + } levels_offset = end_offset; } @@ -713,6 +760,78 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { }) } + /// Writes a chunk in sub-batches sized by the caller-supplied + /// `sub_batch_size` so the post-write data page byte limit check + /// fires before the chunk can grossly overshoot + /// `data_page_size_limit`. + /// + /// For flat (unrepeated) columns sub-batches contain up to + /// `sub_batch_size` levels each. 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. + /// + /// Returns the total number of values consumed across all sub-batches. + /// + /// `#[inline(never)]` keeps this slow path — only reached for + /// variable-width columns whose values need page splitting — out of + /// the hot `write_batch_internal` loop. + #[allow(clippy::too_many_arguments)] + #[inline(never)] + fn write_granular_chunk( + &mut self, + values: &E::Values, + values_offset: usize, + value_indices: Option<&[usize]>, + chunk_size: usize, + chunk_def: LevelDataRef<'_>, + chunk_rep: LevelDataRef<'_>, + sub_batch_size: usize, + ) -> Result { + let mut values_consumed = 0; + let mut sub_start = 0; + while sub_start < 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 + // emitting a `write_mini_batch` per record: records + // average only a handful of levels, so the + // record-at-a-time loop issued roughly `sub_batch_size`× + // more mini-batches than necessary. + let mut e = (sub_start + sub_batch_size).min(chunk_size); + while e < chunk_size && levels[e] != 0 { + e += 1; + } + // `sub_batch_size` can be 0 when the chunker sizes a + // sub-batch below one record; always make progress by + // consuming at least the first whole record. + if e == sub_start { + e = sub_start + 1; + while e < chunk_size && levels[e] != 0 { + e += 1; + } + } + e + } + _ => (sub_start + sub_batch_size).min(chunk_size), + }; + let sub_len = sub_end - sub_start; + let written = self.write_mini_batch( + values, + values_offset + values_consumed, + value_indices, + sub_len, + chunk_def.slice(sub_start, sub_len), + chunk_rep.slice(sub_start, sub_len), + )?; + values_consumed += written; + sub_start = sub_end; + } + Ok(values_consumed) + } + /// 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() { @@ -5211,6 +5330,42 @@ mod tests { } } + #[test] + fn test_level_data_ref_value_count() { + // `value_count` is what the byte-budget chunker uses to convert a + // chunk's level span into a leaf-value count. It must work for any + // column shape — flat, nullable, or nested — because the leaf + // values array is decoupled from the rep/def level stream. + let max_def = 2; + // Non-nullable / unrepeated: no def levels materialized — every + // level is a value. + assert_eq!(LevelDataRef::Absent.value_count(64, max_def), 64); + // Uniform run of present values, and of nulls. + assert_eq!( + LevelDataRef::Uniform { + value: max_def, + count: 40 + } + .value_count(40, max_def), + 40 + ); + assert_eq!( + LevelDataRef::Uniform { + value: max_def - 1, + count: 40 + } + .value_count(40, max_def), + 0 + ); + // Materialized def levels (nullable / nested): only levels equal to + // `max_def` are values; empty-list / null levels are not. + let levels = [2i16, 0, 2, 1, 2, 2, 0]; + assert_eq!( + LevelDataRef::Materialized(&levels).value_count(levels.len(), max_def), + 4 + ); + } + #[test] fn test_uniform_def_levels_all_null() { // All-null column: def_level=0 (null) for every slot, no values written. From 7298c4852df05eaaf9030ff76051df1068bb994a Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Mon, 25 May 2026 12:45:38 -0500 Subject: [PATCH 10/15] fix(parquet): bound page size for arrow byte-array column writes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- parquet/src/arrow/arrow_writer/byte_array.rs | 223 ++++++++++++++++++- 1 file changed, 222 insertions(+), 1 deletion(-) diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index 9cb0718b4d84..c7e172f1a815 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -30,10 +30,12 @@ use crate::geospatial::statistics::GeospatialStatistics; use crate::schema::types::ColumnDescPtr; use crate::util::bit_util::num_required_bits; use crate::util::interner::{Interner, Storage}; +use arrow_array::types::ByteArrayType; use arrow_array::{ Array, ArrayAccessor, BinaryArray, BinaryViewArray, DictionaryArray, FixedSizeBinaryArray, - LargeBinaryArray, LargeStringArray, StringArray, StringViewArray, + GenericByteArray, LargeBinaryArray, LargeStringArray, StringArray, StringViewArray, }; +use arrow_buffer::{ArrowNativeType, Buffer}; use arrow_schema::DataType; macro_rules! downcast_dict_impl { @@ -481,6 +483,90 @@ impl ColumnValueEncoder for ByteArrayEncoder { Ok(()) } + fn count_values_within_byte_budget_gather( + values: &Self::Values, + indices: &[usize], + byte_budget: usize, + ) -> Option { + // `ByteArrayEncoder` only ever writes via `write_gather`, so this + // is the relevant method. + // + // Two-stage walk for the simple offset-buffer byte array types: + // 1. If indices are contiguous, compute the total payload in + // O(1) via a single subtraction on the offsets buffer. + // When the total fits the budget — the overwhelmingly + // common "small values" case — return immediately. + // 2. Otherwise, walk per-value byte sizes from the offsets + // buffer (still cheap, no slice/UTF-8 construction) and + // exit at the first value that pushes the cumulative sum + // past the budget. This bounds skewed distributions: an + // outlier value is caught wherever it lands in the chunk. + let count = match values.data_type() { + DataType::Utf8 => count_within_budget_offsets( + values.as_any().downcast_ref::().unwrap(), + indices, + byte_budget, + ), + DataType::LargeUtf8 => count_within_budget_offsets( + values.as_any().downcast_ref::().unwrap(), + indices, + byte_budget, + ), + DataType::Binary => count_within_budget_offsets( + values.as_any().downcast_ref::().unwrap(), + indices, + byte_budget, + ), + DataType::LargeBinary => count_within_budget_offsets( + values.as_any().downcast_ref::().unwrap(), + indices, + byte_budget, + ), + // View arrays carry each value's length in the low 32 bits of + // its u128 view word, so lengths are scannable without touching + // any data buffer — and the common small-value case skips even + // that scan via an O(1) conservative bound. + DataType::Utf8View => { + let array = values.as_any().downcast_ref::().unwrap(); + count_within_budget_views( + array.views(), + indices, + byte_budget, + max_view_value_len(array.data_buffers()), + ) + } + DataType::BinaryView => { + let array = values.as_any().downcast_ref::().unwrap(); + count_within_budget_views( + array.views(), + indices, + byte_budget, + max_view_value_len(array.data_buffers()), + ) + } + // For arrow Dictionary input, treat every chunk as fitting + // and stay on the batched path. The arrow array being + // Dictionary-encoded in the first place implies its values + // are small enough that dedup is worthwhile, which is the + // opposite of the "5 MiB blob per row" case this fix + // targets. Doing a per-value walk through dict keys (each + // value lookup is keys[i] → values[key] → slice) on every + // chunk costs ~+30-80% vs `main` after writer-dict spill, + // and there is essentially nothing to bound. + DataType::Dictionary(_, _) => indices.len(), + // FixedSizeBinary falls through to the per-value walk via + // `ArrayAccessor::value`. + _ => downcast_op!( + values.data_type(), + values, + count_within_budget_accessor, + indices, + byte_budget + ), + }; + Some(count) + } + fn num_values(&self) -> usize { match &self.dict_encoder { Some(encoder) => encoder.indices.len(), @@ -593,6 +679,141 @@ where } } +/// Cumulative-scan fallback used for byte array types that don't expose +/// a single contiguous offsets buffer — view arrays, dictionary arrays, +/// fixed-size binary. Returns the largest `k` such that the first `k` +/// values picked out by `indices` encode to at most `byte_budget` bytes +/// (or `indices.len()` if they all fit, or `1` if a single value alone +/// exceeds the budget). +/// +/// Free function so it can be used with `downcast_op!`. +fn count_within_budget_accessor(values: T, indices: &[usize], byte_budget: usize) -> usize +where + T: ArrayAccessor + Copy, + T::Item: AsRef<[u8]>, +{ + let mut cum: usize = 0; + for (i, idx) in indices.iter().enumerate() { + let value_len = values.value(*idx).as_ref().len() + std::mem::size_of::(); + cum = cum.saturating_add(value_len); + if cum > byte_budget { + return i.max(1); + } + } + indices.len() +} + +/// Upper bound on any single value's byte length in a view array. +/// +/// An inline view stores at most 12 bytes; an +/// out-of-line view's data is a contiguous slice of exactly one data +/// buffer, so it cannot be longer than the largest data buffer. This is a +/// loose bound (a value is usually far smaller than a whole buffer) but it +/// is O(number of buffers) and always sound. +fn max_view_value_len(buffers: &[Buffer]) -> usize { + /// Bytes that fit inline in a u128 view word (the rest is len + prefix). + const MAX_INLINE_VIEW_LEN: usize = 12; + buffers + .iter() + .map(|b| b.len()) + .max() + .unwrap_or(0) + .max(MAX_INLINE_VIEW_LEN) +} + +/// Two-stage budget count for view arrays (`Utf8View`, `BinaryView`). +/// +/// 1. View arrays have no prefix-sum offsets buffer, so the exact O(1) +/// span subtraction `count_within_budget_offsets` uses is unavailable. +/// But a *conservative* O(1) bound is: every value is at most +/// `max_value_len` bytes, so the whole chunk fits the budget when +/// `n * (max_value_len + 4) <= byte_budget`. This skips the per-value +/// walk for the common small-value case — what view arrays are built +/// for, and exactly the case where there is nothing to bound. +/// 2. Otherwise scan per-value lengths from the low 32 bits of each u128 +/// view word (no data-buffer dereference) and stop at the first value +/// that pushes the cumulative sum past the budget. +fn count_within_budget_views( + views: &[u128], + indices: &[usize], + byte_budget: usize, + max_value_len: usize, +) -> usize { + // Stage 1: O(1) conservative upper bound. + let per_value = max_value_len + std::mem::size_of::(); + if indices.len().saturating_mul(per_value) <= byte_budget { + return indices.len(); + } + // Stage 2: exact per-value scan. + let mut cum: usize = 0; + for (i, idx) in indices.iter().enumerate() { + let len = (views[*idx] as u32) as usize; + cum = cum.saturating_add(len + std::mem::size_of::()); + if cum > byte_budget { + return i.max(1); + } + } + indices.len() +} + +/// Two-stage fast path for `GenericByteArray` +/// (Utf8/LargeUtf8/Binary/LargeBinary). +/// +/// `indices` are assumed sorted ascending — they always are here, since +/// they come from `non_null_indices`, which is built in array order. +/// +/// 1. The span `offsets[last+1] - offsets[first]` is an O(1) upper +/// bound on the chunk's payload: it covers every array position in +/// `[first, last]`, a superset of `indices`. For a non-null chunk +/// `indices` *is* that whole range; for a chunk drawn from a +/// nullable column the skipped positions are nulls, whose offset +/// delta is zero, so the span still equals the exact payload. +/// Either way, if the upper bound fits the budget every value +/// fits — return `indices.len()` with no per-value work. This +/// covers the overwhelmingly common "small values" case for both +/// non-null *and* nullable columns. +/// 2. Otherwise the chunk is genuinely near the budget: walk per-index +/// lengths from the offsets buffer directly (no slice/UTF-8 +/// construction) and stop at the first value that pushes the +/// cumulative sum past the budget. +fn count_within_budget_offsets( + values: &GenericByteArray, + indices: &[usize], + byte_budget: usize, +) -> usize { + if indices.is_empty() { + return 0; + } + let n = indices.len(); + let first = indices[0]; + let last = indices[n - 1]; + let offsets = values.value_offsets(); + let prefix_overhead = std::mem::size_of::(); + + // Stage 1: O(1) span upper bound. Skips Stage 2 in the common case — + // including nullable columns, whose `indices` are sparse. The earlier + // `last - first + 1 == n` contiguity gate forced every nullable chunk + // onto the O(n) Stage 2 walk even though the span check is valid for + // any sorted index set. + if last >= first { + let payload = (offsets[last + 1] - offsets[first]).as_usize(); + if payload + n * prefix_overhead <= byte_budget { + return n; + } + } + + // Stage 2: scan per-index lengths from the offsets buffer. + let mut cum: usize = 0; + for (i, idx) in indices.iter().enumerate() { + let len = (offsets[idx + 1] - offsets[*idx]).as_usize() + prefix_overhead; + cum = cum.saturating_add(len); + if cum > byte_budget { + return i.max(1); + } + } + n +} + /// Computes the min and max for the provided array and indices /// /// This is a free function so it can be used with `downcast_op!` From ef7800b1dbce0164dbb213a0febac100d11a6753 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Mon, 25 May 2026 12:45:47 -0500 Subject: [PATCH 11/15] fix(parquet): bound dictionary page size during dictionary encoding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../src/column/writer/byte_budget_chunker.rs | 67 +++++++++++++------ parquet/tests/arrow_writer_layout.rs | 17 +++-- 2 files changed, 56 insertions(+), 28 deletions(-) diff --git a/parquet/src/column/writer/byte_budget_chunker.rs b/parquet/src/column/writer/byte_budget_chunker.rs index 465568b4caf6..d7276381e655 100644 --- a/parquet/src/column/writer/byte_budget_chunker.rs +++ b/parquet/src/column/writer/byte_budget_chunker.rs @@ -51,6 +51,12 @@ pub(crate) struct ByteBudgetChunker { /// decision short-circuit with no work for every numeric, bool, or /// narrow `FIXED_LEN_BYTE_ARRAY` column. static_always_fits: bool, + /// Configured dictionary page byte limit for the column. + dict_page_byte_limit: usize, + /// As [`Self::static_always_fits`] but for the dictionary page: `true` + /// when one `base_batch_size` mini-batch of this fixed-width type cannot + /// overshoot `dict_page_byte_limit` by more than one mini-batch's worth. + static_dict_always_fits: bool, } impl ByteBudgetChunker { @@ -61,6 +67,7 @@ impl ByteBudgetChunker { base_batch_size: usize, ) -> Self { let page_byte_limit = props.column_data_page_size_limit(descr.path()); + let dict_page_byte_limit = props.column_dictionary_page_size_limit(descr.path()); let static_bytes_per_value = match descr.physical_type() { Type::BOOLEAN => Some(1), Type::INT32 | Type::FLOAT => Some(std::mem::size_of::()), @@ -69,27 +76,36 @@ impl ByteBudgetChunker { Type::FIXED_LEN_BYTE_ARRAY => Some(descr.type_length().max(0) as usize), Type::BYTE_ARRAY => None, }; - let static_always_fits = static_bytes_per_value - .map(|b| b.saturating_mul(base_batch_size) <= page_byte_limit) - .unwrap_or(false); + let static_fits = |limit: usize| { + static_bytes_per_value + .map(|b| b.saturating_mul(base_batch_size) <= limit) + .unwrap_or(false) + }; Self { page_byte_limit, max_def_level: descr.max_def_level(), - static_always_fits, + static_always_fits: static_fits(page_byte_limit), + dict_page_byte_limit, + static_dict_always_fits: static_fits(dict_page_byte_limit), } } /// Decide how many levels at the start of a chunk belong in one - /// mini-batch, so the mini-batch cannot overflow the data page that is - /// currently accumulating value bytes. A returned value smaller than - /// `chunk_size` triggers granular sub-batching in + /// 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`. /// + /// While dictionary-encoding, the data page holds only small RLE indices, + /// but the dictionary page accumulates the distinct values themselves — + /// so it is the dictionary page's remaining budget that must bound the + /// 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, when the column is dictionary-encoding (the data page then - /// holds only small RLE indices, so the plain-encoded byte estimate does - /// not apply), or when the column is a fixed-width type whose - /// mini-batches statically cannot overshoot the data page. + /// 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`. @@ -106,27 +122,34 @@ impl ByteBudgetChunker { if chunk_size == 0 { return chunk_size; } - // While dictionary-encoding, the data page holds only small RLE - // indices — the plain-encoded byte estimate would spuriously shrink - // pages — so stay on the batched fast path. - if encoder.has_dictionary() { - return chunk_size; - } - if self.static_always_fits { - return chunk_size; - } + let budget = if encoder.has_dictionary() { + if self.static_dict_always_fits { + return chunk_size; + } + // 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, + } + } else { + if self.static_always_fits { + return chunk_size; + } + self.page_byte_limit + }; self.byte_budget_sub_batch_size::( values, value_indices, chunk_def, values_offset, chunk_size, - self.page_byte_limit, + budget, ) } /// Inspect value sizes to decide how many of the chunk's values fit in - /// `budget` bytes (the data page remaining budget). + /// `budget` bytes (the data page or dictionary page remaining budget). /// /// `#[inline(never)]` keeps this slow path out of the hot /// `write_batch_internal` loop; numeric and bool columns never reach it. diff --git a/parquet/tests/arrow_writer_layout.rs b/parquet/tests/arrow_writer_layout.rs index b9d997beb289..183590f9b847 100644 --- a/parquet/tests/arrow_writer_layout.rs +++ b/parquet/tests/arrow_writer_layout.rs @@ -408,16 +408,16 @@ fn test_string() { columns: vec![ColumnChunk { pages: vec![ Page { - rows: 130, + rows: 125, page_header_size: 38, - compressed_size: 138, + compressed_size: 114, encoding: Encoding::RLE_DICTIONARY, page_type: PageType::DATA_PAGE, }, Page { - rows: 1250, + rows: 1255, page_header_size: 40, - compressed_size: 10000, + compressed_size: 10040, encoding: Encoding::PLAIN, page_type: PageType::DATA_PAGE, }, @@ -429,10 +429,15 @@ fn test_string() { page_type: PageType::DATA_PAGE, }, ], + // The byte-budget chunker now sub-batches the + // dictionary-encoding phase, so the dictionary page is + // bounded at the 1000-byte limit instead of overshooting + // to 1040 — the dictionary spills one mini-batch earlier + // (125 rows rather than 130). dictionary_page: Some(Page { - rows: 130, + rows: 125, page_header_size: 38, - compressed_size: 1040, + compressed_size: 1000, encoding: Encoding::PLAIN, page_type: PageType::DICTIONARY_PAGE, }), From b3446a5de2d17e7af680bfdd4c2bc1db70415155 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 28 May 2026 11:29:10 -0500 Subject: [PATCH 12/15] fix(parquet): include boundary value in mini-batch byte-budget count 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) --- parquet/src/arrow/arrow_writer/byte_array.rs | 15 ++++++++------- parquet/src/column/writer/encoder.rs | 9 ++++++--- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index c7e172f1a815..1b36a624588e 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -681,10 +681,11 @@ where /// Cumulative-scan fallback used for byte array types that don't expose /// a single contiguous offsets buffer — view arrays, dictionary arrays, -/// fixed-size binary. Returns the largest `k` such that the first `k` -/// values picked out by `indices` encode to at most `byte_budget` bytes -/// (or `indices.len()` if they all fit, or `1` if a single value alone -/// exceeds the budget). +/// fixed-size binary. Returns `indices.len()` if every value fits the +/// budget, otherwise the smallest `k ≥ 1` whose first `k` values' encoded +/// size first exceeds `byte_budget` — i.e. the boundary value is included +/// so the caller's page-flush check trips immediately on this mini-batch, +/// without leaving a sliver to glue onto the next page. /// /// Free function so it can be used with `downcast_op!`. fn count_within_budget_accessor(values: T, indices: &[usize], byte_budget: usize) -> usize @@ -697,7 +698,7 @@ where let value_len = values.value(*idx).as_ref().len() + std::mem::size_of::(); cum = cum.saturating_add(value_len); if cum > byte_budget { - return i.max(1); + return i + 1; } } indices.len() @@ -750,7 +751,7 @@ fn count_within_budget_views( let len = (views[*idx] as u32) as usize; cum = cum.saturating_add(len + std::mem::size_of::()); if cum > byte_budget { - return i.max(1); + return i + 1; } } indices.len() @@ -808,7 +809,7 @@ fn count_within_budget_offsets( let len = (offsets[idx + 1] - offsets[*idx]).as_usize() + prefix_overhead; cum = cum.saturating_add(len); if cum > byte_budget { - return i.max(1); + return i + 1; } } n diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index 83c08e7423fa..33011bd22f5e 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -307,12 +307,15 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { // pushes us past the budget. This both bounds skewed // distributions (one outlier among small values is caught when // it lands, regardless of position) and short-circuits when an - // early value alone exceeds the budget. + // early value alone exceeds the budget. The boundary value is + // included in the returned count so the caller's page-flush check + // trips on this mini-batch rather than leaving a sliver to glue + // onto the next page (see #9972 discussion). let mut cum: usize = 0; for (i, v) in values[start..end].iter().enumerate() { cum = cum.saturating_add(plain_encoded_byte_size::(v)); if cum > byte_budget { - return Some(i.max(1)); + return Some(i + 1); } } Some(n) @@ -334,7 +337,7 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { let Some(v) = values.get(*idx) else { continue }; cum = cum.saturating_add(plain_encoded_byte_size::(v)); if cum > byte_budget { - return Some(i.max(1)); + return Some(i + 1); } } Some(indices.len()) From bcdb878e82b014499bef9e986ece477a95b7387b Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 28 May 2026 11:49:01 -0500 Subject: [PATCH 13/15] test(parquet): update arrow_writer_layout dict-spill expectations 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) --- parquet/tests/arrow_writer_layout.rs | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/parquet/tests/arrow_writer_layout.rs b/parquet/tests/arrow_writer_layout.rs index 183590f9b847..fb4d535a6c98 100644 --- a/parquet/tests/arrow_writer_layout.rs +++ b/parquet/tests/arrow_writer_layout.rs @@ -408,16 +408,16 @@ fn test_string() { columns: vec![ColumnChunk { pages: vec![ Page { - rows: 125, + rows: 126, page_header_size: 38, compressed_size: 114, encoding: Encoding::RLE_DICTIONARY, page_type: PageType::DATA_PAGE, }, Page { - rows: 1255, + rows: 1254, page_header_size: 40, - compressed_size: 10040, + compressed_size: 10032, encoding: Encoding::PLAIN, page_type: PageType::DATA_PAGE, }, @@ -429,15 +429,16 @@ fn test_string() { page_type: PageType::DATA_PAGE, }, ], - // The byte-budget chunker now sub-batches the - // dictionary-encoding phase, so the dictionary page is - // bounded at the 1000-byte limit instead of overshooting - // to 1040 — the dictionary spills one mini-batch earlier - // (125 rows rather than 130). + // The byte-budget chunker sub-batches the dictionary + // phase. The mini-batch deliberately includes the value + // that crosses the 1000-byte limit so the spill triggers + // on this chunk rather than carrying a sliver into the + // next page (see #9972 discussion), giving a 126-row + // dictionary page at 1008 bytes. dictionary_page: Some(Page { - rows: 125, + rows: 126, page_header_size: 38, - compressed_size: 1000, + compressed_size: 1008, encoding: Encoding::PLAIN, page_type: PageType::DICTIONARY_PAGE, }), From 3a7161fab7265f2c586fd39fa478312db28276e0 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 3 Jun 2026 15:39:03 -0500 Subject: [PATCH 14/15] fix(parquet): clean up byte-array page-sizing dead code + lint 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) --- parquet/src/arrow/arrow_writer/byte_array.rs | 41 ++++---------------- parquet/tests/arrow_writer_layout.rs | 10 ++--- 2 files changed, 12 insertions(+), 39 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index 1b36a624588e..01089a6b6f21 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -554,15 +554,13 @@ impl ColumnValueEncoder for ByteArrayEncoder { // chunk costs ~+30-80% vs `main` after writer-dict spill, // and there is essentially nothing to bound. DataType::Dictionary(_, _) => indices.len(), - // FixedSizeBinary falls through to the per-value walk via - // `ArrayAccessor::value`. - _ => downcast_op!( - values.data_type(), - values, - count_within_budget_accessor, - indices, - byte_budget - ), + // Every byte-array type `ByteArrayEncoder` is constructed for + // has an explicit arm above. A `Dictionary(value = FixedSizeBinary)` + // column hits the `Dictionary(_, _)` arm (its `values.data_type()` + // is `Dictionary`), and a bare `FixedSizeBinary` column is routed + // to the generic column writer, never this encoder — so no other + // type can reach here. + data_type => unreachable!("ByteArrayEncoder cannot be constructed for {data_type:?}"), }; Some(count) } @@ -679,31 +677,6 @@ where } } -/// Cumulative-scan fallback used for byte array types that don't expose -/// a single contiguous offsets buffer — view arrays, dictionary arrays, -/// fixed-size binary. Returns `indices.len()` if every value fits the -/// budget, otherwise the smallest `k ≥ 1` whose first `k` values' encoded -/// size first exceeds `byte_budget` — i.e. the boundary value is included -/// so the caller's page-flush check trips immediately on this mini-batch, -/// without leaving a sliver to glue onto the next page. -/// -/// Free function so it can be used with `downcast_op!`. -fn count_within_budget_accessor(values: T, indices: &[usize], byte_budget: usize) -> usize -where - T: ArrayAccessor + Copy, - T::Item: AsRef<[u8]>, -{ - let mut cum: usize = 0; - for (i, idx) in indices.iter().enumerate() { - let value_len = values.value(*idx).as_ref().len() + std::mem::size_of::(); - cum = cum.saturating_add(value_len); - if cum > byte_budget { - return i + 1; - } - } - indices.len() -} - /// Upper bound on any single value's byte length in a view array. /// /// An inline view stores at most 12 bytes; an diff --git a/parquet/tests/arrow_writer_layout.rs b/parquet/tests/arrow_writer_layout.rs index cd1225f64f08..ca1b6a371e65 100644 --- a/parquet/tests/arrow_writer_layout.rs +++ b/parquet/tests/arrow_writer_layout.rs @@ -627,7 +627,8 @@ fn test_fixed_size_binary() { .flat_map(|i| vec![i as u8; value_size]) .collect(); let array = - Arc::new(FixedSizeBinaryArray::try_new(value_size as i32, values.into(), None).unwrap()) as _; + Arc::new(FixedSizeBinaryArray::try_new(value_size as i32, values.into(), None).unwrap()) + as _; let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); let props = WriterProperties::builder() @@ -683,9 +684,9 @@ fn test_dictionary() { // column (one small dictionary page + one RLE_DICTIONARY data page). let num_rows = 2000; let dict_values = StringArray::from_iter_values(["alpha", "beta", "gamma", "delta"]); - let keys = Int32Array::from_iter_values((0..num_rows as i32).map(|i| i % 4)); - let array = Arc::new(DictionaryArray::::try_new(keys, Arc::new(dict_values)).unwrap()) - as _; + let keys = Int32Array::from_iter_values((0..num_rows).map(|i| i % 4)); + let array = + Arc::new(DictionaryArray::::try_new(keys, Arc::new(dict_values)).unwrap()) as _; let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); let props = WriterProperties::builder() @@ -721,4 +722,3 @@ fn test_dictionary() { }, }); } - From 67587ad7242e435c5213d61e12cad2e0cf9e1077 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 3 Jun 2026 16:51:26 -0400 Subject: [PATCH 15/15] test(parquet): trim verbose comments in arrow_writer_layout tests Remove the multi-line explanatory comment blocks from `test_fixed_size_binary` and `test_dictionary`, keeping the one-line summaries. The detail belongs in the PR discussion, not inline in the test bodies. Co-Authored-By: Claude Opus 4.8 (1M context) --- parquet/tests/arrow_writer_layout.rs | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/parquet/tests/arrow_writer_layout.rs b/parquet/tests/arrow_writer_layout.rs index ca1b6a371e65..a604e3f59a5f 100644 --- a/parquet/tests/arrow_writer_layout.rs +++ b/parquet/tests/arrow_writer_layout.rs @@ -611,16 +611,6 @@ fn test_per_column_data_page_size_limit() { #[test] fn test_fixed_size_binary() { // FixedSizeBinary values larger than the data page byte limit. - // - // The arrow writer routes a (non-dictionary) FixedSizeBinary column - // through the *generic* column writer - // (`ColumnValueEncoderImpl`), not the byte-array - // encoder. Its byte-budget chunker uses the plain-encoded value size — - // `type_length` bytes, with no length prefix for FIXED_LEN_BYTE_ARRAY — - // to cap each data page. With 1 KiB values and a 4 KiB page limit, a - // mini-batch grows until its cumulative bytes first cross the limit - // (the boundary value is included), so 5 values land per page (5120 B) - // instead of buffering all 64 values into one ~64 KiB page. let value_size = 1024usize; let num_rows = 64usize; let values: Vec = (0..num_rows) @@ -672,16 +662,6 @@ fn test_fixed_size_binary() { #[test] fn test_dictionary() { // Arrow `DictionaryArray` input. - // - // The byte-budget chunker's gather path has a dedicated arm for - // dictionary-typed arrow input that reports "everything fits" - // (`indices.len()`), keeping the column on the batched fast path: an - // arrow array that is *already* dictionary-encoded implies its values - // are small enough that dedup is worthwhile — the opposite of the large - // blob case the page-size fix targets — so there is nothing to bound and - // walking dict keys per chunk would only cost time. This test drives that - // arm and confirms the column still lays out as a normal dictionary - // column (one small dictionary page + one RLE_DICTIONARY data page). let num_rows = 2000; let dict_values = StringArray::from_iter_values(["alpha", "beta", "gamma", "delta"]); let keys = Int32Array::from_iter_values((0..num_rows).map(|i| i % 4));