This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11798-69be95c1d925bd876c71c11ec6cb2a646d54a1ad in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit 4910a816d3f63312a660cab696bdfb3c0b60406d Author: Thomas Tanon <[email protected]> AuthorDate: Mon Oct 26 17:54:57 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;
