laskoviymishka commented on code in PR #2741:
URL: https://github.com/apache/iceberg-rust/pull/2741#discussion_r3966648935
##########
crates/iceberg/src/writer/file_writer/parquet_writer.rs:
##########
@@ -75,9 +76,34 @@ impl ParquetWriterBuilder {
props,
schema,
match_mode,
+ row_group_size_bytes: None,
}
}
+ /// Target an on-disk **encoded (post-compression)** size for each row
group.
Review Comment:
This says the target is an on-disk, post-compression size "regardless of
compression ratio", but I don't think `in_progress_size()` gives us that.
parquet-rs documents it as the anticipated *encoded* size of the in-progress
group — the buffered, pre-file-compression estimate. With snappy/zstd the
on-disk groups come out a fraction of the target (a 64 MiB target lands ~10-25
MB groups), so "regardless of compression ratio" is the opposite of what
happens. The test only uses `UNCOMPRESSED`, so the claim never actually gets
exercised.
I'd reword to "estimated encoded (pre-compression) size" and note that with
a codec enabled the on-disk groups are smaller in proportion to the compression
ratio. Same wording is in the comment on the flush block in `write()` (around
line 584) — worth fixing both. wdyt?
##########
crates/iceberg/src/writer/file_writer/parquet_writer.rs:
##########
@@ -548,6 +579,24 @@ impl FileWriter for ParquetWriter {
.with_source(err)
})?;
+ // Cut a row group once its anticipated encoded size reaches the
configured
+ // byte target. `in_progress_size` is parquet's estimate of the
post-encoding
+ // (post-compression) size of the open row group, so this bounds row
groups by
+ // on-disk bytes regardless of compression ratio. The row-count cap in
+ // `writer_properties` still applies — whichever fires first cuts the
group.
+ if self
+ .row_group_size_bytes
+ .is_some_and(|target| writer.in_progress_size() >= target)
Review Comment:
With `bytes == 0` this is always true, so we'd flush after every write and
emit one row group per batch (and possibly an empty group right after a flush).
I'd guard it — either `assert!(bytes > 0, ...)` in the builder or return a
`Result` with `ErrorKind::DataInvalid`. wdyt?
##########
crates/iceberg/src/writer/file_writer/parquet_writer.rs:
##########
@@ -75,9 +76,34 @@ impl ParquetWriterBuilder {
props,
schema,
match_mode,
+ row_group_size_bytes: None,
}
}
+ /// Target an on-disk **encoded (post-compression)** size for each row
group.
+ ///
+ /// [`WriterProperties::max_row_group_size`] bounds a row group only by row
+ /// *count*, so the on-disk size of a group depends on how well the data
+ /// compresses — uniform-row-count groups can vary wildly in bytes. When
this
+ /// is set, the writer additionally cuts a row group as soon as its
anticipated
+ /// encoded size (parquet's [`AsyncArrowWriter::in_progress_size`]) reaches
+ /// `bytes`, yielding byte-uniform row groups without the caller having to
+ /// predict a compression ratio. This is the analogue of Java Iceberg's
+ /// `write.parquet.row-group-size-bytes`.
+ ///
+ /// The row-count cap still applies as an upper bound; whichever limit is
hit
+ /// first cuts the group.
+ ///
+ /// The size is checked **after each [`FileWriter::write`] call**
(parquet's
+ /// recommended `in_progress_size`-driven flush pattern), so a row group
can
+ /// overshoot the target by at most one written batch. Callers that hand
the
+ /// writer one very large batch will therefore see coarser cuts; streaming
+ /// callers that write in modest batches get byte-uniform row groups.
+ pub fn with_row_group_size_bytes(mut self, bytes: usize) -> Self {
Review Comment:
`from_table_properties` already maps `write.parquet.row-group-size-bytes`
onto `WriterProperties::set_max_row_group_bytes`, and the async writer honours
that natively — it splits batches at the boundary and doesn't overshoot. So we
end up with two paths to the same intent: `from_table_properties` gets the
native cut, `new().with_row_group_size_bytes(N)` gets this post-write check
with up to a full batch of overshoot. A caller can even chain both and get two
conflicting limits.
Could we forward to the native property instead of carrying a separate field
and flush?
```rust
pub fn with_row_group_size_bytes(mut self, bytes: usize) -> Self {
self.props = self.props.into_builder()
.set_max_row_group_bytes(Some(bytes))
.build();
self
}
```
That drops `row_group_size_bytes` and the flush block entirely and makes
both entry points behave identically. If there's a reason the post-write check
is preferable here that I'm missing, I'd love the docstring to spell it out
rather than claim equivalence to the Java property. wdyt?
##########
crates/iceberg/src/writer/file_writer/parquet_writer.rs:
##########
@@ -896,6 +945,107 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn test_parquet_writer_row_group_size_bytes() -> Result<()> {
+ let temp_dir = TempDir::new().unwrap();
+ let file_io = FileIO::new_with_fs();
+ let location_gen = DefaultLocationGenerator::with_data_location(
+ temp_dir.path().to_str().unwrap().to_string(),
+ );
+ let file_name_gen =
+ DefaultFileNameGenerator::new("test".to_string(), None,
DataFileFormat::Parquet);
+
+ let schema = {
+ let fields =
+ vec![
+ Field::new("col", DataType::Int64,
false).with_metadata(HashMap::from([(
+ PARQUET_FIELD_ID_META_KEY.to_string(),
+ "0".to_string(),
+ )])),
+ ];
+ Arc::new(arrow_schema::Schema::new(fields))
+ };
+ // Pseudo-random, ~incompressible values with dictionary + compression
off, so
+ // the encoded size is ~8 bytes/row and a small byte target
deterministically
+ // cuts many row groups regardless of encoding heuristics.
+ let values: Vec<i64> = (0..100_000i64)
+ .map(|i| (i as u64).wrapping_mul(2_654_435_761) as i64)
+ .collect();
+
+ let output_file = file_io.new_output(
+ location_gen.generate_location(None,
&file_name_gen.generate_file_name()),
+ )?;
+
+ // ~800 KB of encoded data with a 64 KiB row-group byte target. The
row-count
+ // cap is left at the parquet default (~1M rows), so only the byte
target can
+ // cut the data into multiple groups. The size is checked after each
write, so
+ // feed the rows in modest batches (as a streaming caller would)
rather than
+ // one giant batch.
+ let target_bytes: usize = 64 * 1024;
+ let mut pw = ParquetWriterBuilder::new(
+ WriterProperties::builder()
+ .set_dictionary_enabled(false)
+ .set_compression(parquet::basic::Compression::UNCOMPRESSED)
+ .build(),
+ Arc::new(schema.as_ref().try_into().unwrap()),
+ )
+ .with_row_group_size_bytes(target_bytes)
+ .build(output_file)
+ .await?;
+ for chunk in values.chunks(4_000) {
+ let col =
Arc::new(Int64Array::from_iter_values(chunk.iter().copied())) as ArrayRef;
+ let batch = RecordBatch::try_new(schema.clone(),
vec![col]).unwrap();
+ pw.write(&batch).await?;
+ }
+ let res = pw.close().await?;
+ assert_eq!(res.len(), 1);
+ let data_file = res
+ .into_iter()
+ .next()
+ .unwrap()
+ .content(DataContentType::Data)
+ .partition(Struct::empty())
+ .partition_spec_id(0)
+ .build()
+ .unwrap();
+
+ assert_eq!(data_file.record_count(), 100_000);
+
+ // The byte target is far below the total encoded size, so the
size-based cut
+ // fired repeatedly instead of emitting one default row group. ~800 KB
/ 64 KiB
+ // with a one-batch (32 KiB) overshoot tolerance lands at ~9 groups;
assert a
+ // band rather than an exact count so a parquet patch bump nudging the
+ // in_progress_size estimate doesn't make this brittle.
+ let input_content = file_io
+ .new_input(data_file.file_path.clone())?
+ .read()
+ .await?;
+ let metadata =
+
parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder::try_new(input_content)
+ .unwrap()
+ .metadata()
+ .clone();
+ let row_group_sizes: Vec<i64> = (0..metadata.num_row_groups())
+ .map(|i| metadata.row_group(i).compressed_size())
+ .collect();
+ assert!(
+ (5..=15).contains(&row_group_sizes.len()),
+ "expected ~9 byte-targeted row groups, got {}",
+ row_group_sizes.len()
+ );
+ // Every group but the last (the smaller remainder) crossed the target
before
+ // being cut and overshoots by at most one written batch, so each sits
in a
+ // band around the 64 KiB target — the byte-sizing contract.
+ for &size in &row_group_sizes[..row_group_sizes.len() - 1] {
+ assert!(
+ (target_bytes as i64 / 2..=target_bytes as i64 *
3).contains(&size),
+ "row group size {size} B is outside the byte-target band
(target {target_bytes} B)"
Review Comment:
The comment says a non-final group overshoots by at most one written batch
(~96 KiB), but the asserted ceiling here is `target * 3` = 192 KiB — loose
enough that a regression firing the cut two batches late would still pass. I'd
tighten the upper bound to `target_bytes + one_batch_bytes` so the band
actually pins the contract the comment describes.
##########
crates/iceberg/src/writer/file_writer/parquet_writer.rs:
##########
@@ -548,6 +579,24 @@ impl FileWriter for ParquetWriter {
.with_source(err)
})?;
+ // Cut a row group once its anticipated encoded size reaches the
configured
+ // byte target. `in_progress_size` is parquet's estimate of the
post-encoding
+ // (post-compression) size of the open row group, so this bounds row
groups by
+ // on-disk bytes regardless of compression ratio. The row-count cap in
+ // `writer_properties` still applies — whichever fires first cuts the
group.
+ if self
+ .row_group_size_bytes
+ .is_some_and(|target| writer.in_progress_size() >= target)
+ {
+ writer.flush().await.map_err(|err| {
+ Error::new(
+ ErrorKind::Unexpected,
+ "Failed to flush row group in parquet writer.",
Review Comment:
tiny nit while we're here — the sibling write error reads "Failed to write
using parquet writer."; this one says "in parquet writer". I'd match them
("...using parquet writer.").
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
