This is an automated email from the ASF dual-hosted git repository.
Jefffrey 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 cd7c6b83ab fix(parquet): split row groups iteratively instead of
recursively (#10626)
cd7c6b83ab is described below
commit cd7c6b83abd6605a83014b3d043930a592542510
Author: Stefan Wang <[email protected]>
AuthorDate: Sat Aug 22 06:44:07 2026 -0700
fix(parquet): split row groups iteratively instead of recursively (#10626)
# Which issue does this PR close?
- Closes #9386.
# Rationale for this change
`ArrowWriter::write` splits a batch by calling itself on each half, so
it recurses once per row group. A row limit much smaller than the batch
turns that into one stack frame per row group, and the process aborts:
```
thread '...' has overflowed its stack
fatal runtime error: stack overflow, aborting
```
The limits are caller-supplied, so a large batch written under a small
`max_row_group_row_count` takes down the process rather than returning
an error.
# What changes are included in this PR?
`write` now loops over the rows still to be written instead of
recursing. Each pass fills the in-progress row group, flushes when a
limit is reached, and carries the remainder into the next pass, so stack
use is constant no matter how many row groups a batch produces.
The split points and flush conditions are unchanged.
# Are these changes tested?
Yes. `test_row_group_limit_rows_only_many_splits` writes 50,000 rows
with `max_row_group_row_count = 1`, so the batch splits into 50,000 row
groups. On the current code it aborts the process:
```console
$ cargo test -p parquet --lib test_row_group_limit_rows_only_many_splits
running 1 test
thread
'arrow::arrow_writer::tests::test_row_group_limit_rows_only_many_splits'
(19274774) has overflowed its stack
fatal runtime error: stack overflow, aborting
error: test failed, to rerun pass `-p parquet --lib`
Caused by:
process didn't exit successfully: `.../parquet-80a96ba37d90b1a2
test_row_group_limit_rows_only_many_splits` (signal: 6, SIGABRT: process abort
signal)
```
`test_row_group_limit_both_apply_to_same_batch` pins the split points
where both limits bite the same batch: the row limit trims it to 5 rows
and the byte limit then trims those to 4, giving row groups of `[14,
6]`. It passes on the current code, so it holds the iterative version to
the same split points.
Both, plus the rest of the `test_row_group_limit_*` cases, on this
branch:
<details><summary>Raw output</summary>
```console
$ cargo test -p parquet --lib test_row_group_limit
test
arrow::arrow_writer::tests::test_row_group_limit_none_writes_single_row_group
... ok
test arrow::arrow_writer::tests::test_row_group_limit_rows_only ... ok
test
arrow::arrow_writer::tests::test_row_group_limit_both_apply_to_same_batch ... ok
test
arrow::arrow_writer::tests::test_row_group_limit_both_row_wins_multiple_batches
... ok
test
arrow::arrow_writer::tests::test_row_group_limit_both_row_wins_single_batch ...
ok
test arrow::arrow_writer::tests::test_row_group_limit_bytes_only ... ok
test
arrow::arrow_writer::tests::test_row_group_limit_bytes_flushes_when_current_group_already_too_large
... ok
test arrow::arrow_writer::tests::test_row_group_limit_both_bytes_wins ... ok
test arrow::arrow_writer::tests::test_row_group_limit_rows_only_many_splits
... ok
test result: ok. 9 passed; 0 failed; 0 ignored; 0 measured; 1252 filtered
out; finished in 1.20s
$ cargo test -p parquet --lib
test result: ok. 1261 passed; 0 failed; 0 ignored; 0 measured; 0 filtered
out; finished in 13.29s
```
</details>
The byte limit reaches its split path only once rows are already
buffered, and re-enters with an empty row group, so it stays shallow —
the row count limit was the one that could run away.
# Are there any user-facing changes?
No API change. Writes that previously aborted the process now complete.
---------
Signed-off-by: 1fanwang <[email protected]>
---
parquet/src/arrow/arrow_writer/mod.rs | 185 +++++++++++++++++++++++-----------
1 file changed, 128 insertions(+), 57 deletions(-)
diff --git a/parquet/src/arrow/arrow_writer/mod.rs
b/parquet/src/arrow/arrow_writer/mod.rs
index c49c919fa0..f8acb65044 100644
--- a/parquet/src/arrow/arrow_writer/mod.rs
+++ b/parquet/src/arrow/arrow_writer/mod.rs
@@ -362,74 +362,92 @@ impl<W: Write + Send> ArrowWriter<W> {
return Ok(());
}
- let in_progress = match &mut self.in_progress {
- Some(in_progress) => in_progress,
- x => x.insert(
- self.row_group_writer_factory
-
.create_row_group_writer(self.writer.flushed_row_groups().len())?,
- ),
- };
-
- if let Some(max_rows) = self.max_row_group_row_count
- && in_progress.buffered_rows + batch.num_rows() > max_rows
- {
- let to_write = max_rows - in_progress.buffered_rows;
- let a = batch.slice(0, to_write);
- let b = batch.slice(to_write, batch.num_rows() - to_write);
- self.write(&a)?;
- return self.write(&b);
- }
+ // Rows not yet handed to a row group writer. Splitting iterates here
instead of
+ // recursing, so a small row group limit over a large batch cannot
exhaust the stack.
+ let mut remaining = batch.clone();
+
+ loop {
+ let in_progress = match &mut self.in_progress {
+ Some(in_progress) => in_progress,
+ x => x.insert(
+ self.row_group_writer_factory
+
.create_row_group_writer(self.writer.flushed_row_groups().len())?,
+ ),
+ };
+ let buffered_rows = in_progress.buffered_rows;
- // Check byte limit: if we have buffered data, use measured average
row size
- // to split batch proactively before exceeding byte limit
- if let Some(max_bytes) = self.max_row_group_bytes
- && in_progress.buffered_rows > 0
- {
- let current_bytes = in_progress.get_estimated_total_bytes();
+ // Leading rows of `remaining` that still fit in the current row
group, when the
+ // rest has to go to a later one.
+ let mut split_at = match self.max_row_group_row_count {
+ Some(max_rows) if buffered_rows + remaining.num_rows() >
max_rows => {
+ Some(max_rows - buffered_rows)
+ }
+ _ => None,
+ };
- if current_bytes >= max_bytes {
- self.flush()?;
- return self.write(batch);
- }
+ // Check byte limit: if we have buffered data, use measured
average row size
+ // to split batch proactively before exceeding byte limit. Both
limits apply to
+ // the same rows, so measure against whatever the row limit
already trimmed
+ // `remaining` down to; otherwise the row limit would always win.
+ let candidate_rows = split_at.unwrap_or_else(||
remaining.num_rows());
- if let Some(avg_row_bytes) = current_bytes
- .checked_div(in_progress.buffered_rows)
- .filter(|avg_row_bytes| *avg_row_bytes > 0)
+ if let Some(max_bytes) = self.max_row_group_bytes
+ && buffered_rows > 0
{
- // At this point, `current_bytes < max_bytes` (checked above)
- let remaining_bytes = max_bytes - current_bytes;
- let rows_that_fit =
remaining_bytes.checked_div(avg_row_bytes).unwrap_or(0);
-
- if batch.num_rows() > rows_that_fit {
- if rows_that_fit > 0 {
- let a = batch.slice(0, rows_that_fit);
- let b = batch.slice(rows_that_fit, batch.num_rows() -
rows_that_fit);
- self.write(&a)?;
- return self.write(&b);
- } else {
- self.flush()?;
- return self.write(batch);
+ let current_bytes = in_progress.get_estimated_total_bytes();
+
+ if current_bytes >= max_bytes {
+ self.flush()?;
+ continue;
+ }
+
+ if let Some(avg_row_bytes) = current_bytes
+ .checked_div(buffered_rows)
+ .filter(|avg_row_bytes| *avg_row_bytes > 0)
+ {
+ // At this point, `current_bytes < max_bytes` (checked
above)
+ let remaining_bytes = max_bytes - current_bytes;
+ let rows_that_fit =
remaining_bytes.checked_div(avg_row_bytes).unwrap_or(0);
+
+ if candidate_rows > rows_that_fit {
+ if rows_that_fit > 0 {
+ split_at = Some(rows_that_fit);
+ } else {
+ self.flush()?;
+ continue;
+ }
}
}
}
- }
- match self.cdc_chunkers.as_mut() {
- Some(chunkers) => in_progress.write_with_chunkers(batch,
chunkers)?,
- None => in_progress.write(batch)?,
- }
+ let rest = split_at.map(|to_write| {
+ let rest = remaining.slice(to_write, remaining.num_rows() -
to_write);
+ remaining = remaining.slice(0, to_write);
+ rest
+ });
- let should_flush = self
- .max_row_group_row_count
- .is_some_and(|max| in_progress.buffered_rows >= max)
- || self
- .max_row_group_bytes
- .is_some_and(|max| in_progress.get_estimated_total_bytes() >=
max);
+ let in_progress = self.in_progress.as_mut().unwrap();
+ match self.cdc_chunkers.as_mut() {
+ Some(chunkers) => in_progress.write_with_chunkers(&remaining,
chunkers)?,
+ None => in_progress.write(&remaining)?,
+ }
+
+ let should_flush = self
+ .max_row_group_row_count
+ .is_some_and(|max| in_progress.buffered_rows >= max)
+ || self
+ .max_row_group_bytes
+ .is_some_and(|max| in_progress.get_estimated_total_bytes()
>= max);
+
+ if should_flush {
+ self.flush()?
+ }
- if should_flush {
- self.flush()?
+ match rest {
+ Some(rest) => remaining = rest,
+ None => return Ok(()),
+ }
}
- Ok(())
}
/// Writes the given buf bytes to the internal buffer.
@@ -5788,6 +5806,34 @@ mod tests {
);
}
+ #[test]
+ // A row limit far smaller than the batch splits it many times over; the
split must not
+ // consume stack proportional to the number of row groups.
+ fn test_row_group_limit_rows_only_many_splits() {
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(Some(1))
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let rows = 50_000;
+ let builder = write_batches(
+ WriteBatchesShape {
+ num_batches: 1,
+ rows_per_batch: rows,
+ row_size: 4,
+ },
+ props,
+ );
+
+ let sizes = row_group_sizes(builder.metadata());
+ assert_eq!(sizes.len(), rows, "Every row should get its own row
group");
+ assert_eq!(
+ sizes.iter().sum::<i64>(),
+ rows as i64,
+ "Total rows should be preserved"
+ );
+ }
+
#[test]
// When only max_row_group_bytes is set, respect the byte limit
fn test_row_group_limit_bytes_only() {
@@ -5944,6 +5990,31 @@ mod tests {
assert_eq!(total_rows, 100, "Total rows should be preserved");
}
+ #[test]
+ // Both limits can apply to the same batch: the row limit trims it to 5
rows, and the
+ // byte limit then trims those 5 down to 4.
+ fn test_row_group_limit_both_apply_to_same_batch() {
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(Some(15))
+ .set_max_row_group_bytes(Some(1500))
+ .build();
+
+ let builder = write_batches(
+ WriteBatchesShape {
+ num_batches: 2,
+ rows_per_batch: 10,
+ row_size: 100,
+ },
+ props,
+ );
+
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[14, 6],
+ "Byte limit should still apply to a batch the row limit already
split"
+ );
+ }
+
#[test]
fn arrow_column_chunk_close_mut_drops_column_index() {
use crate::arrow::ArrowSchemaConverter;
