This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 91c5db296f fix(parquet): apply max_row_group_bytes to the first batch
of a row group (#11798)
91c5db296f is described below
commit 91c5db296f8477d384abdc427301ab9e8baac537
Author: Thomas Tanon <[email protected]>
AuthorDate: Mon Oct 26 18:36:55 2026 +0000
fix(parquet): apply max_row_group_bytes to the first batch of a row group
(#11798)
# Which issue does this PR close?
- Closes #11792
- Closes https://github.com/apache/arrow-rs/pull/11796
# Rationale for this change
`max_row_group_bytes` has no effect when a single batch is larger than
the limit: the first batch written to a row group is accepted whole, as
there are no buffered rows to estimate a row size from. Row groups then
hold one batch each, however large it is. See the issue for a
reproduction.
# What changes are included in this PR?
When a byte limit is set and the in-progress row group is empty,
`ArrowWriter::write` now writes a single row first. That gives the
existing logic a measured row size to work out how many of the remaining
rows fit, so they get split the same way as for a non-empty row group.
The row limit still applies first, and nothing changes when
`max_row_group_bytes` is unset.
# Are these changes tested?
Yes, `test_row_group_limit_bytes_splits_first_batch` writes 100-row
batches against a byte limit of about 35 rows. It checks that no row
group holds a whole batch and that every row group stays within one row
of the limit. It fails on `main` with row groups of `[100, 100, 100]`.
The existing `test_row_group_limit_*` tests pass unchanged.
# Are there any significant user-facing changes?
Writers that set `max_row_group_bytes` with batches larger than it now
get row groups that respect the limit, so they produce more and smaller
row groups than before. The first row of each row group is encoded on
its own. That adds a small per-row-group cost, but only when a byte
limit is set.
Disclaimer: mostly vibe coded but carefully reviewed
Co-authored-by: RIchard Baah <[email protected]>
---
parquet/src/arrow/arrow_writer/mod.rs | 48 ++++++++++++++++++++++++++++++++---
1 file changed, 45 insertions(+), 3 deletions(-)
diff --git a/parquet/src/arrow/arrow_writer/mod.rs
b/parquet/src/arrow/arrow_writer/mod.rs
index 8ce51de52d..4b03e72f42 100644
--- a/parquet/src/arrow/arrow_writer/mod.rs
+++ b/parquet/src/arrow/arrow_writer/mod.rs
@@ -393,9 +393,14 @@ impl<W: Write + Send> ArrowWriter<W> {
// `remaining` down to; otherwise the row limit would always win.
let candidate_rows = split_at.unwrap_or_else(||
remaining.num_rows());
- if let Some(max_bytes) = self.max_row_group_bytes
- && buffered_rows > 0
- {
+ if self.max_row_group_bytes.is_some() && buffered_rows == 0 {
+ // There is no row size to extrapolate from yet: write a
single row first,
+ // so that the byte limit also applies to the rest of the
batch instead of
+ // letting a batch larger than the limit fill the row group on
its own.
+ if candidate_rows > 1 {
+ split_at = Some(1);
+ }
+ } else if let Some(max_bytes) = self.max_row_group_bytes {
let current_bytes = in_progress.get_estimated_total_bytes();
if current_bytes >= max_bytes {
@@ -4983,6 +4988,43 @@ mod tests {
);
}
+ #[test]
+ // A batch far larger than the byte limit must be split even when it is
the first one
+ // written to a row group, where there are no buffered rows to estimate
the row size from.
+ fn test_row_group_limit_bytes_splits_first_batch() {
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(Some(3500))
+ .build();
+
+ let builder = write_batches(
+ WriteBatchesShape {
+ num_batches: 3,
+ rows_per_batch: 100,
+ row_size: 100,
+ },
+ props,
+ );
+
+ let sizes = row_group_sizes(builder.metadata());
+ assert!(
+ sizes.iter().all(|&s| s < 100),
+ "No row group should hold a whole batch, got {sizes:?}",
+ );
+ assert_eq!(
+ sizes.iter().sum::<i64>(),
+ 300,
+ "Total rows should be preserved"
+ );
+ for row_group in builder.metadata().row_groups() {
+ assert!(
+ row_group.total_byte_size() <= 3500 + 100,
+ "Row groups should stay within one row of the byte limit, got
{}",
+ row_group.total_byte_size(),
+ );
+ }
+ }
+
#[test]
fn arrow_column_chunk_close_mut_drops_column_index() {
use crate::arrow::ArrowSchemaConverter;
