Jefffrey commented on code in PR #10626:
URL: https://github.com/apache/arrow-rs/pull/10626#discussion_r3800244656
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -360,74 +360,93 @@ 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())?,
- ),
- };
+ // 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 {
+ if self.in_progress.is_none() {
+ let row_group_index = self.writer.flushed_row_groups().len();
+ self.in_progress = Some(
+ self.row_group_writer_factory
+ .create_row_group_writer(row_group_index)?,
+ );
+ }
+ let buffered_rows =
self.in_progress.as_ref().unwrap().buffered_rows;
- 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);
- }
+ // 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,
+ };
- // 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();
+ // Check byte limit: if we have buffered data, use measured
average row size
+ // to split batch proactively before exceeding byte limit
+ if split_at.is_none()
Review Comment:
codex helped me identify a possible deviation: by checking `split_at` here
it means the row count limit can win over the byte limit. for example this test
will pass on main with recursive version:
```rust
#[test]
fn test123() {
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,
);
let sizes = row_group_sizes(builder.metadata());
assert_eq!(sizes, &[14, 6]);
}
```
but with this iterative version it fails because the sizes becomes `[15, 5]`
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -360,74 +360,93 @@ 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())?,
- ),
- };
+ // 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 {
+ if self.in_progress.is_none() {
+ let row_group_index = self.writer.flushed_row_groups().len();
+ self.in_progress = Some(
+ self.row_group_writer_factory
+ .create_row_group_writer(row_group_index)?,
+ );
+ }
+ let buffered_rows =
self.in_progress.as_ref().unwrap().buffered_rows;
- 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);
- }
+ // 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,
+ };
- // 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();
+ // Check byte limit: if we have buffered data, use measured
average row size
+ // to split batch proactively before exceeding byte limit
+ if split_at.is_none()
+ && let Some(max_bytes) = self.max_row_group_bytes
+ && buffered_rows > 0
+ {
+ let current_bytes = self
+ .in_progress
+ .as_ref()
+ .unwrap()
+ .get_estimated_total_bytes();
- if current_bytes >= max_bytes {
- self.flush()?;
- return self.write(batch);
- }
+ if current_bytes >= max_bytes {
+ self.flush()?;
+ continue;
+ }
- if let Some(avg_row_bytes) = current_bytes
- .checked_div(in_progress.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 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);
+ 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 remaining.num_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 in_progress = self.in_progress.as_mut().unwrap();
Review Comment:
it would be nice if we could align with the previous version in not needing
to do these unwraps multiple times
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
