You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "yonipeleg33 (via GitHub)" <gi...@apache.org> on 2026/02/05 09:33:21 UTC
[PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
yonipeleg33 opened a new pull request, #9357:
URL: https://github.com/apache/arrow-rs/pull/9357
# Which issue does this PR close?
This PR implements another suggestion introduced in https://github.com/apache/arrow-rs/issues/1213:
> Add functionality to flush based on a bytes threshold instead of, or in addition to, the current row threshold
So not "Closes" anything new.
# Rationale for this change
A best effort to match Spark's (or more specifically, Hadoop's) `parquet.block.size` configuration behaviour, as documented in parquet-hadoop's [README](https://github.com/apache/parquet-java/blob/master/parquet-hadoop/README.md):
> **Property:** `parquet.block.size`
> **Description:** The block size in bytes...
Since arrow's parquet writer writes batches, it's inherently different than Hadoop's per-record writer behaviour - so the behaviour of `max_row_group_bytes` will be different than Hadoop's `parquet.block.size`, but this is the closest I could reasonably get (see details below).
# What changes are included in this PR?
**Configuration changes**
- New optional `max_row_group_bytes` configuration option in `WriterProperties`
- Rename existing `max_row_group_size` private property to `max_row_group_row_count`
- **Backwards compatible:** No public APIs changed:
- `set_max_row_group_size()` and `max_row_group_size()` still remain with their existing signatures.
- added `set_max_row_group_row_count()` and `max_row_group_row_count()` which expose the `Option<usize>` type.
- If `set_max_row_group_row_count(None)` is called, `max_row_group_size()` will return `usize::MAX`.
**Writer changes**
`ArrowWriter::write` now supports any combination of these two properties (row count and row bytes):
- Both are unset -> Write everything in one row group.
- One is set -> Respect only this one (either bytes or rows amount).
- Both are set -> Respect the lower of them: Open a new row group when either row count or byte size limits reached
Byte limit is calculated once per batch (as opposed to Hadoop's per-record calculation):
Before writing each batch, compute the average row size in bytes based on previous writes, and flush or split the batch according to that average before hitting the limit.
This means that the first batch will always be written as a whole (unless row count limit is also set).
# Are these changes tested?
Yes - added unit tests to check all different combinations of these two properties being set.
# Are there any user-facing changes?
Yes: Introducing new APIs to configure byte limits on row groups, and slight change to existing one (returning `usize::MAX` from `max_row_group_size()` if it was unset by the user).
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "yonipeleg33 (via GitHub)" <gi...@apache.org>.
yonipeleg33 commented on PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#issuecomment-3873666595
Thanks so much for the review, guys!
I'll wait for @alamb to take one more look as well.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#issuecomment-3881058994
I made a small follow on PR to add some additional tests
- https://github.com/apache/arrow-rs/pull/9387
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769392681
##########
parquet/src/file/properties.rs:
##########
@@ -575,7 +595,34 @@ impl WriterPropertiesBuilder {
/// If the value is set to 0.
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
+ self.max_row_group_row_count = Some(value);
+ self
+ }
+
+ /// Sets maximum number of rows in a row group, or `None` for unlimited.
+ ///
+ /// # Panics
+ /// If the value is `Some(0)`.
+ pub fn set_max_row_group_row_count(mut self, value: Option<usize>) -> Self {
+ if let Some(v) = value {
+ assert!(v > 0, "Cannot have a 0 max row group size");
+ }
+ self.max_row_group_row_count = value;
+ self
+ }
+
+ /// Sets maximum size of a row group in bytes, or `None` for unlimited.
+ ///
+ /// Row groups are flushed when their estimated encoded size exceeds this threshold.
+ /// This is similar to parquet-mr's `parquet.block.size` behavior.
Review Comment:
parquet-mr is just the official java implementation for parquet, you can rewrite the comment to clarify that this match the official parquet Jqva implementation
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "yonipeleg33 (via GitHub)" <gi...@apache.org>.
yonipeleg33 commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2784353097
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -331,18 +341,57 @@ impl<W: Write + Send> ArrowWriter<W> {
),
};
- // If would exceed max_row_group_size, split batch
- if in_progress.buffered_rows + batch.num_rows() > self.max_row_group_size {
- let to_write = self.max_row_group_size - 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);
+ if let Some(max_rows) = self.max_row_group_row_count {
+ if 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);
+ }
+ }
+
+ // 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 {
+ if in_progress.buffered_rows > 0 {
+ let current_bytes = in_progress.get_estimated_total_bytes();
+
+ if current_bytes >= max_bytes {
+ self.flush()?;
+ return self.write(batch);
+ }
+
+ let avg_row_bytes = current_bytes / in_progress.buffered_rows;
+ if avg_row_bytes > 0 {
+ let remaining_bytes = max_bytes - current_bytes;
Review Comment:
Good point. Done
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769659979
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
+ RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
+ }
+
+ #[test]
+ fn test_row_group_limit_none_writes_single_row_group() {
+ // When both limits are None, all data should go into a single row group
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[1000],
+ "With no limits, all rows should be in a single row group"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_rows_only() {
+ // When only max_row_group_size is set, respect the row limit
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(300)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[300, 300, 300, 100],
+ "Row groups should be split by row count"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_bytes_only() {
+ // When only max_row_group_bytes is set, respect the byte limit
+ // Create batches with string data for more predictable byte sizes
+ // Write in multiple small batches so byte-based splitting can work
+ // (first batch establishes the avg row size, subsequent batches are split)
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "str",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+
+ // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(Some(3500))
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
+
+ // Write 10 batches of 10 rows each (100 rows total)
+ // Each string is ~100 bytes
+ for batch_idx in 0..10 {
+ let strings: Vec<String> = (0..10)
+ .map(|i| format!("{:0>100}", batch_idx * 10 + i))
+ .collect();
+ let array = StringArray::from(strings);
+ let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
+ writer.write(&batch).unwrap();
+ }
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ let sizes = row_group_sizes(builder.metadata());
+
+ assert!(
+ sizes.len() > 1,
+ "Should have multiple row groups due to byte limit, got {sizes:?}",
+ );
+
+ let total_rows: i64 = sizes.iter().sum();
+ assert_eq!(total_rows, 100, "Total rows should be preserved");
+ }
+
+ #[test]
+ fn test_row_group_limit_both_row_wins() {
+ // When both limits are set, the row limit triggers first
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(200) // Will trigger at 200 rows
+ .set_max_row_group_bytes(Some(1024 * 1024)) // 1MB - won't trigger for small int data
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[200, 200, 200, 200, 200],
+ "Row limit should trigger before byte limit"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_both_bytes_wins() {
+ // When both limits are set, the byte limit triggers first
+ // Write in multiple small batches so byte-based splitting can work
Review Comment:
According to the comment on the method, the way you write batches should not affect, only the config, that is if the byte based got hit first, it should write that, if the row hit first it should write that
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "yonipeleg33 (via GitHub)" <gi...@apache.org>.
yonipeleg33 commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2771005709
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
+ RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
+ }
+
+ #[test]
+ fn test_row_group_limit_none_writes_single_row_group() {
+ // When both limits are None, all data should go into a single row group
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[1000],
+ "With no limits, all rows should be in a single row group"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_rows_only() {
+ // When only max_row_group_size is set, respect the row limit
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(300)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[300, 300, 300, 100],
+ "Row groups should be split by row count"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_bytes_only() {
+ // When only max_row_group_bytes is set, respect the byte limit
+ // Create batches with string data for more predictable byte sizes
+ // Write in multiple small batches so byte-based splitting can work
+ // (first batch establishes the avg row size, subsequent batches are split)
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "str",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+
+ // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(Some(3500))
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
+
+ // Write 10 batches of 10 rows each (100 rows total)
+ // Each string is ~100 bytes
+ for batch_idx in 0..10 {
+ let strings: Vec<String> = (0..10)
+ .map(|i| format!("{:0>100}", batch_idx * 10 + i))
+ .collect();
+ let array = StringArray::from(strings);
+ let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
+ writer.write(&batch).unwrap();
+ }
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ let sizes = row_group_sizes(builder.metadata());
+
+ assert!(
+ sizes.len() > 1,
+ "Should have multiple row groups due to byte limit, got {sizes:?}",
+ );
+
+ let total_rows: i64 = sizes.iter().sum();
+ assert_eq!(total_rows, 100, "Total rows should be preserved");
+ }
+
+ #[test]
+ fn test_row_group_limit_both_row_wins() {
Review Comment:
AFAICT, that beats the purpose of this configuration: Its purpose is to control the IO profile of the writer (i.e. how much and when it writes to disk), and for that, the data needs to _at least_ be already encoded before calculating the row group size.
This is also backed by the Java source code:
it [calculates](https://github.com/apache/parquet-java/blob/7190ab6a571c3cffd31e267042c193e90c0301ad/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/InternalParquetRecordWriter.java#L183) `memSize` using `columnStore.getBufferedSize()`, which is [documented](https://github.com/apache/parquet-java/blob/7190ab6a571c3cffd31e267042c193e90c0301ad/parquet-column/src/main/java/org/apache/parquet/column/ColumnWriteStore.java#L54) as follows:
> `@return approximate size of the buffered encoded binary data`
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2777513088
##########
parquet/src/file/properties.rs:
##########
@@ -575,7 +595,34 @@ impl WriterPropertiesBuilder {
/// If the value is set to 0.
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
+ self.max_row_group_row_count = Some(value);
+ self
+ }
+
+ /// Sets maximum number of rows in a row group, or `None` for unlimited.
+ ///
+ /// # Panics
+ /// If the value is `Some(0)`.
+ pub fn set_max_row_group_row_count(mut self, value: Option<usize>) -> Self {
+ if let Some(v) = value {
+ assert!(v > 0, "Cannot have a 0 max row group size");
+ }
+ self.max_row_group_row_count = value;
+ self
+ }
+
+ /// Sets maximum size of a row group in bytes, or `None` for unlimited.
+ ///
+ /// Row groups are flushed when their estimated encoded size exceeds this threshold.
+ /// This is similar to parquet-mr's `parquet.block.size` behavior.
Review Comment:
Also, parquet-mr is I think now officially called "parquet-java" https://github.com/apache/parquet-java
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769448374
##########
parquet/src/file/properties.rs:
##########
@@ -575,7 +595,34 @@ impl WriterPropertiesBuilder {
/// If the value is set to 0.
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
Review Comment:
Should we deprecate this function function?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769572896
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -187,8 +187,11 @@ pub struct ArrowWriter<W: Write> {
/// Creates new [`ArrowRowGroupWriter`] instances as required
row_group_writer_factory: ArrowRowGroupWriterFactory,
- /// The length of arrays to write to each row group
- max_row_group_size: usize,
+ /// The maximum number of rows to write to each row group, or None for unlimited
+ max_row_group_row_count: Option<usize>,
+
+ /// The maximum size in bytes for a row group, or None for unlimited
+ max_row_group_bytes: Option<usize>,
Review Comment:
can they both be None?
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -187,8 +187,11 @@ pub struct ArrowWriter<W: Write> {
/// Creates new [`ArrowRowGroupWriter`] instances as required
row_group_writer_factory: ArrowRowGroupWriterFactory,
- /// The length of arrays to write to each row group
- max_row_group_size: usize,
+ /// The maximum number of rows to write to each row group, or None for unlimited
+ max_row_group_row_count: Option<usize>,
+
+ /// The maximum size in bytes for a row group, or None for unlimited
+ max_row_group_bytes: Option<usize>,
Review Comment:
can they both be None?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "yonipeleg33 (via GitHub)" <gi...@apache.org>.
yonipeleg33 commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769983367
##########
parquet/src/file/properties.rs:
##########
@@ -575,7 +595,34 @@ impl WriterPropertiesBuilder {
/// If the value is set to 0.
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
+ self.max_row_group_row_count = Some(value);
+ self
+ }
+
+ /// Sets maximum number of rows in a row group, or `None` for unlimited.
+ ///
+ /// # Panics
+ /// If the value is `Some(0)`.
Review Comment:
Done
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
+ RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
+ }
+
+ #[test]
+ fn test_row_group_limit_none_writes_single_row_group() {
+ // When both limits are None, all data should go into a single row group
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[1000],
+ "With no limits, all rows should be in a single row group"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_rows_only() {
+ // When only max_row_group_size is set, respect the row limit
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(300)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[300, 300, 300, 100],
+ "Row groups should be split by row count"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_bytes_only() {
+ // When only max_row_group_bytes is set, respect the byte limit
+ // Create batches with string data for more predictable byte sizes
+ // Write in multiple small batches so byte-based splitting can work
+ // (first batch establishes the avg row size, subsequent batches are split)
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "str",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+
+ // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
Review Comment:
Done
##########
parquet/src/file/properties.rs:
##########
@@ -45,6 +45,8 @@ pub const DEFAULT_STATISTICS_ENABLED: EnabledStatistics = EnabledStatistics::Pag
pub const DEFAULT_WRITE_PAGE_HEADER_STATISTICS: bool = false;
/// Default value for [`WriterProperties::max_row_group_size`]
pub const DEFAULT_MAX_ROW_GROUP_SIZE: usize = 1024 * 1024;
+/// Default value for [`WriterProperties::max_row_group_bytes`] (128 MB, same as parquet-mr's parquet.block.size)
Review Comment:
Done
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
+ RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
+ }
+
+ #[test]
+ fn test_row_group_limit_none_writes_single_row_group() {
+ // When both limits are None, all data should go into a single row group
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[1000],
+ "With no limits, all rows should be in a single row group"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_rows_only() {
+ // When only max_row_group_size is set, respect the row limit
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(300)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[300, 300, 300, 100],
+ "Row groups should be split by row count"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_bytes_only() {
+ // When only max_row_group_bytes is set, respect the byte limit
+ // Create batches with string data for more predictable byte sizes
+ // Write in multiple small batches so byte-based splitting can work
+ // (first batch establishes the avg row size, subsequent batches are split)
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "str",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+
+ // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(Some(3500))
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
+
+ // Write 10 batches of 10 rows each (100 rows total)
+ // Each string is ~100 bytes
+ for batch_idx in 0..10 {
+ let strings: Vec<String> = (0..10)
+ .map(|i| format!("{:0>100}", batch_idx * 10 + i))
+ .collect();
+ let array = StringArray::from(strings);
+ let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
+ writer.write(&batch).unwrap();
+ }
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ let sizes = row_group_sizes(builder.metadata());
+
+ assert!(
+ sizes.len() > 1,
+ "Should have multiple row groups due to byte limit, got {sizes:?}",
+ );
+
+ let total_rows: i64 = sizes.iter().sum();
+ assert_eq!(total_rows, 100, "Total rows should be preserved");
+ }
+
+ #[test]
+ fn test_row_group_limit_both_row_wins() {
Review Comment:
No can do; The first batch is always written as a whole, because we need some statistics in order to calculate average row size. This is also noted in the PR description:
> This means that the first batch will always be written as a whole (unless row count limit is also set).
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
+ RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
+ }
+
+ #[test]
+ fn test_row_group_limit_none_writes_single_row_group() {
+ // When both limits are None, all data should go into a single row group
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[1000],
+ "With no limits, all rows should be in a single row group"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_rows_only() {
+ // When only max_row_group_size is set, respect the row limit
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(300)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[300, 300, 300, 100],
+ "Row groups should be split by row count"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_bytes_only() {
+ // When only max_row_group_bytes is set, respect the byte limit
+ // Create batches with string data for more predictable byte sizes
+ // Write in multiple small batches so byte-based splitting can work
+ // (first batch establishes the avg row size, subsequent batches are split)
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "str",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+
+ // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(Some(3500))
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
+
+ // Write 10 batches of 10 rows each (100 rows total)
+ // Each string is ~100 bytes
+ for batch_idx in 0..10 {
+ let strings: Vec<String> = (0..10)
+ .map(|i| format!("{:0>100}", batch_idx * 10 + i))
+ .collect();
+ let array = StringArray::from(strings);
+ let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
+ writer.write(&batch).unwrap();
+ }
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ let sizes = row_group_sizes(builder.metadata());
+
+ assert!(
+ sizes.len() > 1,
+ "Should have multiple row groups due to byte limit, got {sizes:?}",
+ );
+
+ let total_rows: i64 = sizes.iter().sum();
+ assert_eq!(total_rows, 100, "Total rows should be preserved");
+ }
+
+ #[test]
+ fn test_row_group_limit_both_row_wins() {
+ // When both limits are set, the row limit triggers first
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(200) // Will trigger at 200 rows
+ .set_max_row_group_bytes(Some(1024 * 1024)) // 1MB - won't trigger for small int data
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[200, 200, 200, 200, 200],
+ "Row limit should trigger before byte limit"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_both_bytes_wins() {
+ // When both limits are set, the byte limit triggers first
+ // Write in multiple small batches so byte-based splitting can work
Review Comment:
Unfortunately, the way data is fed _does_ affect the row group splits, because of the first batch issue (noted in the PR description):
> This means that the first batch will always be written as a whole (unless row count limit is also set).
And even beyond the first batch, the behaviour is not predictable: Byte-based limit is enforced by calculating the average row size, based on previous batches. This is more dynamic than row-based limit.
I'm not sure what's actionable from this comment. If you think there's still a missing test case, please LMK.
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
Review Comment:
Done
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
+ RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
+ }
+
+ #[test]
+ fn test_row_group_limit_none_writes_single_row_group() {
+ // When both limits are None, all data should go into a single row group
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[1000],
+ "With no limits, all rows should be in a single row group"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_rows_only() {
+ // When only max_row_group_size is set, respect the row limit
Review Comment:
Done
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
Review Comment:
Nice! Done
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
+ RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
+ }
+
+ #[test]
+ fn test_row_group_limit_none_writes_single_row_group() {
+ // When both limits are None, all data should go into a single row group
Review Comment:
Done (moved to above the test function)
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
+ RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
+ }
+
+ #[test]
+ fn test_row_group_limit_none_writes_single_row_group() {
+ // When both limits are None, all data should go into a single row group
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[1000],
+ "With no limits, all rows should be in a single row group"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_rows_only() {
+ // When only max_row_group_size is set, respect the row limit
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(300)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[300, 300, 300, 100],
+ "Row groups should be split by row count"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_bytes_only() {
+ // When only max_row_group_bytes is set, respect the byte limit
+ // Create batches with string data for more predictable byte sizes
+ // Write in multiple small batches so byte-based splitting can work
+ // (first batch establishes the avg row size, subsequent batches are split)
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "str",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+
+ // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(Some(3500))
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
+
+ // Write 10 batches of 10 rows each (100 rows total)
+ // Each string is ~100 bytes
+ for batch_idx in 0..10 {
+ let strings: Vec<String> = (0..10)
+ .map(|i| format!("{:0>100}", batch_idx * 10 + i))
+ .collect();
+ let array = StringArray::from(strings);
+ let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
+ writer.write(&batch).unwrap();
+ }
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ let sizes = row_group_sizes(builder.metadata());
+
+ assert!(
+ sizes.len() > 1,
+ "Should have multiple row groups due to byte limit, got {sizes:?}",
+ );
+
+ let total_rows: i64 = sizes.iter().sum();
+ assert_eq!(total_rows, 100, "Total rows should be preserved");
+ }
+
+ #[test]
+ fn test_row_group_limit_both_row_wins() {
+ // When both limits are set, the row limit triggers first
Review Comment:
Done
##########
parquet/src/file/properties.rs:
##########
@@ -575,7 +595,34 @@ impl WriterPropertiesBuilder {
/// If the value is set to 0.
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
Review Comment:
Probably. Done
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
+ RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
+ }
+
+ #[test]
+ fn test_row_group_limit_none_writes_single_row_group() {
+ // When both limits are None, all data should go into a single row group
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[1000],
+ "With no limits, all rows should be in a single row group"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_rows_only() {
+ // When only max_row_group_size is set, respect the row limit
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(300)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[300, 300, 300, 100],
+ "Row groups should be split by row count"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_bytes_only() {
+ // When only max_row_group_bytes is set, respect the byte limit
+ // Create batches with string data for more predictable byte sizes
+ // Write in multiple small batches so byte-based splitting can work
+ // (first batch establishes the avg row size, subsequent batches are split)
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "str",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+
+ // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(Some(3500))
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
+
+ // Write 10 batches of 10 rows each (100 rows total)
+ // Each string is ~100 bytes
+ for batch_idx in 0..10 {
+ let strings: Vec<String> = (0..10)
+ .map(|i| format!("{:0>100}", batch_idx * 10 + i))
+ .collect();
+ let array = StringArray::from(strings);
+ let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
+ writer.write(&batch).unwrap();
+ }
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ let sizes = row_group_sizes(builder.metadata());
+
+ assert!(
+ sizes.len() > 1,
+ "Should have multiple row groups due to byte limit, got {sizes:?}",
+ );
+
+ let total_rows: i64 = sizes.iter().sum();
+ assert_eq!(total_rows, 100, "Total rows should be preserved");
+ }
+
+ #[test]
+ fn test_row_group_limit_both_row_wins() {
+ // When both limits are set, the row limit triggers first
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(200) // Will trigger at 200 rows
+ .set_max_row_group_bytes(Some(1024 * 1024)) // 1MB - won't trigger for small int data
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[200, 200, 200, 200, 200],
+ "Row limit should trigger before byte limit"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_both_bytes_wins() {
Review Comment:
Done - see `test_row_group_limit_both_row_wins_multiple_batches` vs. `test_row_group_limit_both_row_wins_single_batch`
##########
parquet/src/file/properties.rs:
##########
@@ -575,7 +595,34 @@ impl WriterPropertiesBuilder {
/// If the value is set to 0.
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
+ self.max_row_group_row_count = Some(value);
+ self
+ }
+
+ /// Sets maximum number of rows in a row group, or `None` for unlimited.
+ ///
+ /// # Panics
+ /// If the value is `Some(0)`.
+ pub fn set_max_row_group_row_count(mut self, value: Option<usize>) -> Self {
+ if let Some(v) = value {
+ assert!(v > 0, "Cannot have a 0 max row group size");
+ }
+ self.max_row_group_row_count = value;
+ self
+ }
+
+ /// Sets maximum size of a row group in bytes, or `None` for unlimited.
+ ///
+ /// Row groups are flushed when their estimated encoded size exceeds this threshold.
+ /// This is similar to the official `parquet.block.size` behavior.
+ ///
+ /// # Panics
+ /// If the value is `Some(0)`.
+ pub fn set_max_row_group_bytes(mut self, value: Option<usize>) -> Self {
+ if let Some(v) = value {
+ assert!(v > 0, "Cannot have a 0 max row group bytes");
+ }
Review Comment:
Neat, done
##########
parquet/src/file/properties.rs:
##########
@@ -575,7 +595,34 @@ impl WriterPropertiesBuilder {
/// If the value is set to 0.
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
+ self.max_row_group_row_count = Some(value);
+ self
+ }
+
+ /// Sets maximum number of rows in a row group, or `None` for unlimited.
+ ///
+ /// # Panics
+ /// If the value is `Some(0)`.
+ pub fn set_max_row_group_row_count(mut self, value: Option<usize>) -> Self {
+ if let Some(v) = value {
+ assert!(v > 0, "Cannot have a 0 max row group size");
+ }
+ self.max_row_group_row_count = value;
+ self
+ }
+
+ /// Sets maximum size of a row group in bytes, or `None` for unlimited.
+ ///
+ /// Row groups are flushed when their estimated encoded size exceeds this threshold.
+ /// This is similar to the official `parquet.block.size` behavior.
Review Comment:
Done
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2790427629
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -331,18 +341,58 @@ impl<W: Write + Send> ArrowWriter<W> {
),
};
- // If would exceed max_row_group_size, split batch
- if in_progress.buffered_rows + batch.num_rows() > self.max_row_group_size {
- let to_write = self.max_row_group_size - 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);
+ if let Some(max_rows) = self.max_row_group_row_count {
+ if 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);
Review Comment:
here is a reproducer (I will file a follow on ticket)
```rust
#[test]
fn test_row_group_limit_rows_only_pathological_stack_overflow_demo() {
let schema = Arc::new(Schema::new(vec![Field::new(
"int",
ArrowDataType::Int32,
false,
)]));
let array = Int32Array::from((0..1_000_000_i32).collect::<Vec<_>>());
let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
let props = WriterProperties::builder()
.set_max_row_group_row_count(Some(1))
.set_max_row_group_bytes(None)
.build();
let file = tempfile::tempfile().unwrap();
let mut writer = ArrowWriter::try_new(file, schema, Some(props)).unwrap();
// This currently recurses once per row-group split and can overflow the stack.
writer.write(&batch).unwrap();
}
```
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "etseidl (via GitHub)" <gi...@apache.org>.
etseidl commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2776350777
##########
parquet/src/file/properties.rs:
##########
@@ -573,9 +594,42 @@ impl WriterPropertiesBuilder {
///
/// # Panics
/// If the value is set to 0.
+ #[deprecated(
+ since = "57.3.0",
+ note = "Use `set_max_row_group_row_count` instead",
+ )]
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
+ self.max_row_group_row_count = Some(value);
+ self
+ }
+
+ /// Sets maximum number of rows in a row group, or `None` for unlimited.
+ ///
+ /// If both `max_row_group_row_count` and `max_row_group_bytes` are set,
+ /// the row group with the smallest limit will be applied.
+ ///
+ /// # Panics
+ /// If the value is `Some(0)`.
+ pub fn set_max_row_group_row_count(mut self, value: Option<usize>) -> Self {
+ assert_ne!(value, Some(0), "Cannot have a 0 max row group bytes");
+ self.max_row_group_row_count = value;
+ self
+ }
+
+ /// Sets maximum size of a row group in bytes, or `None` for unlimited.
+ ///
+ /// Row groups are flushed when their estimated encoded size exceeds this threshold.
+ /// This is similar to the official Java implementation for `parquet.block.size`'s behavior.
+ ///
+ /// If both `max_row_group_row_count` and `max_row_group_bytes` are set,
+ /// the row group with the smallest limit will be applied.
Review Comment:
```suggestion
/// the row group with the smaller limit will be produced.
```
##########
parquet/src/file/properties.rs:
##########
@@ -573,9 +594,42 @@ impl WriterPropertiesBuilder {
///
/// # Panics
/// If the value is set to 0.
+ #[deprecated(
+ since = "57.3.0",
+ note = "Use `set_max_row_group_row_count` instead",
+ )]
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
+ self.max_row_group_row_count = Some(value);
+ self
+ }
+
+ /// Sets maximum number of rows in a row group, or `None` for unlimited.
+ ///
+ /// If both `max_row_group_row_count` and `max_row_group_bytes` are set,
+ /// the row group with the smallest limit will be applied.
Review Comment:
```suggestion
/// the row group with the smaller limit will be produced.
```
##########
parquet/src/file/properties.rs:
##########
@@ -573,9 +594,42 @@ impl WriterPropertiesBuilder {
///
/// # Panics
/// If the value is set to 0.
+ #[deprecated(
+ since = "57.3.0",
+ note = "Use `set_max_row_group_row_count` instead",
+ )]
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
+ self.max_row_group_row_count = Some(value);
+ self
+ }
+
+ /// Sets maximum number of rows in a row group, or `None` for unlimited.
+ ///
+ /// If both `max_row_group_row_count` and `max_row_group_bytes` are set,
+ /// the row group with the smallest limit will be applied.
+ ///
+ /// # Panics
+ /// If the value is `Some(0)`.
+ pub fn set_max_row_group_row_count(mut self, value: Option<usize>) -> Self {
+ assert_ne!(value, Some(0), "Cannot have a 0 max row group bytes");
Review Comment:
```suggestion
assert_ne!(value, Some(0), "Cannot have a 0 max row group row count");
```
##########
parquet/src/file/properties.rs:
##########
@@ -573,9 +594,42 @@ impl WriterPropertiesBuilder {
///
/// # Panics
/// If the value is set to 0.
+ #[deprecated(
+ since = "57.3.0",
Review Comment:
```suggestion
since = "58.0.0",
```
`57.3.0` has already been released
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -314,8 +320,12 @@ impl<W: Write + Send> ArrowWriter<W> {
/// Encodes the provided [`RecordBatch`]
///
/// If this would cause the current row group to exceed [`WriterProperties::max_row_group_size`]
Review Comment:
```suggestion
/// If this would cause the current row group to exceed [`WriterProperties::max_row_group_row_count`]
```
?
##########
parquet/src/file/properties.rs:
##########
@@ -45,6 +45,9 @@ pub const DEFAULT_STATISTICS_ENABLED: EnabledStatistics = EnabledStatistics::Pag
pub const DEFAULT_WRITE_PAGE_HEADER_STATISTICS: bool = false;
/// Default value for [`WriterProperties::max_row_group_size`]
pub const DEFAULT_MAX_ROW_GROUP_SIZE: usize = 1024 * 1024;
+/// Default value for [`WriterProperties::max_row_group_bytes`] (128 MB, same as the official Java
+/// implementation for `parquet.block.size`)
+pub const DEFAULT_MAX_ROW_GROUP_BYTES: usize = 128 * 1024 * 1024;
Review Comment:
This constant appears to be unused. I'd vote for less clutter and get rid of it
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "yonipeleg33 (via GitHub)" <gi...@apache.org>.
yonipeleg33 commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2771029203
##########
parquet/src/file/properties.rs:
##########
@@ -575,7 +595,34 @@ impl WriterPropertiesBuilder {
/// If the value is set to 0.
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
Review Comment:
I do agree that this API should be deprecated. Thanks for pointing it out!
> this is a breaking change, as clippy will fail for users
This is inherently _not_ a breaking change; the purpose of marking APIs as deprecated is to warn users _before_ making a breaking change, without actually making this change.
This PR already calls for a minor bump due to the new APIs introduced; deprecating the old one does not change the version semantics for this PR.
> it might be in a different pr
I don't mind either way: Leaving it here or opening a new PR for deprecating the old API. LMK what's your preference and I'll do it.
> if you keep it please update the pr description under changes to users
Done!
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "yonipeleg33 (via GitHub)" <gi...@apache.org>.
yonipeleg33 commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2778930325
##########
parquet/src/file/properties.rs:
##########
@@ -573,9 +594,42 @@ impl WriterPropertiesBuilder {
///
/// # Panics
/// If the value is set to 0.
+ #[deprecated(
+ since = "57.3.0",
Review Comment:
Thanks, done
##########
parquet/src/file/properties.rs:
##########
@@ -573,9 +594,42 @@ impl WriterPropertiesBuilder {
///
/// # Panics
/// If the value is set to 0.
+ #[deprecated(
+ since = "57.3.0",
+ note = "Use `set_max_row_group_row_count` instead",
+ )]
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
+ self.max_row_group_row_count = Some(value);
+ self
+ }
+
+ /// Sets maximum number of rows in a row group, or `None` for unlimited.
+ ///
+ /// If both `max_row_group_row_count` and `max_row_group_bytes` are set,
+ /// the row group with the smallest limit will be applied.
+ ///
+ /// # Panics
+ /// If the value is `Some(0)`.
+ pub fn set_max_row_group_row_count(mut self, value: Option<usize>) -> Self {
+ assert_ne!(value, Some(0), "Cannot have a 0 max row group bytes");
+ self.max_row_group_row_count = value;
+ self
+ }
+
+ /// Sets maximum size of a row group in bytes, or `None` for unlimited.
+ ///
+ /// Row groups are flushed when their estimated encoded size exceeds this threshold.
+ /// This is similar to the official Java implementation for `parquet.block.size`'s behavior.
+ ///
+ /// If both `max_row_group_row_count` and `max_row_group_bytes` are set,
+ /// the row group with the smallest limit will be applied.
Review Comment:
Done
##########
parquet/src/file/properties.rs:
##########
@@ -45,6 +45,9 @@ pub const DEFAULT_STATISTICS_ENABLED: EnabledStatistics = EnabledStatistics::Pag
pub const DEFAULT_WRITE_PAGE_HEADER_STATISTICS: bool = false;
/// Default value for [`WriterProperties::max_row_group_size`]
pub const DEFAULT_MAX_ROW_GROUP_SIZE: usize = 1024 * 1024;
+/// Default value for [`WriterProperties::max_row_group_bytes`] (128 MB, same as the official Java
+/// implementation for `parquet.block.size`)
+pub const DEFAULT_MAX_ROW_GROUP_BYTES: usize = 128 * 1024 * 1024;
Review Comment:
Removed. Good catch! (Leftover from previous implementations)
##########
parquet/src/file/properties.rs:
##########
@@ -573,9 +594,42 @@ impl WriterPropertiesBuilder {
///
/// # Panics
/// If the value is set to 0.
+ #[deprecated(
+ since = "57.3.0",
+ note = "Use `set_max_row_group_row_count` instead",
+ )]
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
+ self.max_row_group_row_count = Some(value);
+ self
+ }
+
+ /// Sets maximum number of rows in a row group, or `None` for unlimited.
+ ///
+ /// If both `max_row_group_row_count` and `max_row_group_bytes` are set,
+ /// the row group with the smallest limit will be applied.
+ ///
+ /// # Panics
+ /// If the value is `Some(0)`.
+ pub fn set_max_row_group_row_count(mut self, value: Option<usize>) -> Self {
+ assert_ne!(value, Some(0), "Cannot have a 0 max row group bytes");
Review Comment:
Done
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -314,8 +320,12 @@ impl<W: Write + Send> ArrowWriter<W> {
/// Encodes the provided [`RecordBatch`]
///
/// If this would cause the current row group to exceed [`WriterProperties::max_row_group_size`]
Review Comment:
Indeed! Done
##########
parquet/src/file/properties.rs:
##########
@@ -573,9 +594,42 @@ impl WriterPropertiesBuilder {
///
/// # Panics
/// If the value is set to 0.
+ #[deprecated(
+ since = "57.3.0",
+ note = "Use `set_max_row_group_row_count` instead",
+ )]
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
+ self.max_row_group_row_count = Some(value);
+ self
+ }
+
+ /// Sets maximum number of rows in a row group, or `None` for unlimited.
+ ///
+ /// If both `max_row_group_row_count` and `max_row_group_bytes` are set,
+ /// the row group with the smallest limit will be applied.
Review Comment:
Done
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769479288
##########
parquet/src/file/properties.rs:
##########
@@ -45,6 +45,8 @@ pub const DEFAULT_STATISTICS_ENABLED: EnabledStatistics = EnabledStatistics::Pag
pub const DEFAULT_WRITE_PAGE_HEADER_STATISTICS: bool = false;
/// Default value for [`WriterProperties::max_row_group_size`]
pub const DEFAULT_MAX_ROW_GROUP_SIZE: usize = 1024 * 1024;
+/// Default value for [`WriterProperties::max_row_group_bytes`] (128 MB, same as parquet-mr's parquet.block.size)
Review Comment:
same as my other parquet-mr comment
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769653185
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
+ RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
+ }
+
+ #[test]
+ fn test_row_group_limit_none_writes_single_row_group() {
+ // When both limits are None, all data should go into a single row group
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[1000],
+ "With no limits, all rows should be in a single row group"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_rows_only() {
+ // When only max_row_group_size is set, respect the row limit
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(300)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[300, 300, 300, 100],
+ "Row groups should be split by row count"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_bytes_only() {
+ // When only max_row_group_bytes is set, respect the byte limit
+ // Create batches with string data for more predictable byte sizes
+ // Write in multiple small batches so byte-based splitting can work
+ // (first batch establishes the avg row size, subsequent batches are split)
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "str",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+
+ // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(Some(3500))
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
+
+ // Write 10 batches of 10 rows each (100 rows total)
+ // Each string is ~100 bytes
+ for batch_idx in 0..10 {
+ let strings: Vec<String> = (0..10)
+ .map(|i| format!("{:0>100}", batch_idx * 10 + i))
+ .collect();
+ let array = StringArray::from(strings);
+ let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
+ writer.write(&batch).unwrap();
+ }
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ let sizes = row_group_sizes(builder.metadata());
+
+ assert!(
+ sizes.len() > 1,
+ "Should have multiple row groups due to byte limit, got {sizes:?}",
+ );
+
+ let total_rows: i64 = sizes.iter().sum();
+ assert_eq!(total_rows, 100, "Total rows should be preserved");
+ }
+
+ #[test]
+ fn test_row_group_limit_both_row_wins() {
+ // When both limits are set, the row limit triggers first
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(200) // Will trigger at 200 rows
+ .set_max_row_group_bytes(Some(1024 * 1024)) // 1MB - won't trigger for small int data
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[200, 200, 200, 200, 200],
+ "Row limit should trigger before byte limit"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_both_bytes_wins() {
Review Comment:
Can you have a similar test with rows wins that have the same structure but only config change?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769670528
##########
parquet/src/file/properties.rs:
##########
@@ -575,7 +595,34 @@ impl WriterPropertiesBuilder {
/// If the value is set to 0.
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
+ self.max_row_group_row_count = Some(value);
+ self
+ }
+
+ /// Sets maximum number of rows in a row group, or `None` for unlimited.
+ ///
+ /// # Panics
+ /// If the value is `Some(0)`.
+ pub fn set_max_row_group_row_count(mut self, value: Option<usize>) -> Self {
+ if let Some(v) = value {
+ assert!(v > 0, "Cannot have a 0 max row group size");
+ }
+ self.max_row_group_row_count = value;
+ self
+ }
+
+ /// Sets maximum size of a row group in bytes, or `None` for unlimited.
+ ///
+ /// Row groups are flushed when their estimated encoded size exceeds this threshold.
+ /// This is similar to the official `parquet.block.size` behavior.
Review Comment:
```suggestion
/// This is similar to the official Java implementation `parquet.block.size` behavior.
```
this is not part of the spec so there is no official about it
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769452671
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
Review Comment:
```suggestion
/// Each row is 4 bytes (one i32).
```
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
Review Comment:
```suggestion
/// Each row is 4 bytes (one `i32`).
```
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2777513839
##########
parquet/src/file/properties.rs:
##########
@@ -45,6 +45,9 @@ pub const DEFAULT_STATISTICS_ENABLED: EnabledStatistics = EnabledStatistics::Pag
pub const DEFAULT_WRITE_PAGE_HEADER_STATISTICS: bool = false;
/// Default value for [`WriterProperties::max_row_group_size`]
pub const DEFAULT_MAX_ROW_GROUP_SIZE: usize = 1024 * 1024;
+/// Default value for [`WriterProperties::max_row_group_bytes`] (128 MB, same as the official Java
+/// implementation for `parquet.block.size`)
+pub const DEFAULT_MAX_ROW_GROUP_BYTES: usize = 128 * 1024 * 1024;
Review Comment:
Or should we set it?
##########
parquet/src/file/properties.rs:
##########
@@ -247,11 +251,25 @@ impl WriterProperties {
self.write_batch_size
}
- /// Returns maximum number of rows in a row group.
+ /// Returns maximum number of rows in a row group, or `usize::MAX` if unlimited.
///
/// For more details see [`WriterPropertiesBuilder::set_max_row_group_size`]
pub fn max_row_group_size(&self) -> usize {
Review Comment:
Given the introduction of max_row_group_count, what would you think about deprecating `max_row_group_size` and directing people to that new setting?
##########
parquet/src/file/properties.rs:
##########
@@ -468,7 +487,8 @@ impl Default for WriterPropertiesBuilder {
data_page_size_limit: DEFAULT_PAGE_SIZE,
data_page_row_count_limit: DEFAULT_DATA_PAGE_ROW_COUNT_LIMIT,
write_batch_size: DEFAULT_WRITE_BATCH_SIZE,
- max_row_group_size: DEFAULT_MAX_ROW_GROUP_SIZE,
+ max_row_group_row_count: Some(DEFAULT_MAX_ROW_GROUP_SIZE),
Review Comment:
Could we also please align the constant name to the parameter name (eg. `DEFAULT_MAX_ROW_GROUP_COUNT`)
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "yonipeleg33 (via GitHub)" <gi...@apache.org>.
yonipeleg33 commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2779221476
##########
parquet/src/file/properties.rs:
##########
@@ -468,7 +487,8 @@ impl Default for WriterPropertiesBuilder {
data_page_size_limit: DEFAULT_PAGE_SIZE,
data_page_row_count_limit: DEFAULT_DATA_PAGE_ROW_COUNT_LIMIT,
write_batch_size: DEFAULT_WRITE_BATCH_SIZE,
- max_row_group_size: DEFAULT_MAX_ROW_GROUP_SIZE,
+ max_row_group_row_count: Some(DEFAULT_MAX_ROW_GROUP_SIZE),
Review Comment:
Done
##########
parquet/src/file/properties.rs:
##########
@@ -247,11 +251,25 @@ impl WriterProperties {
self.write_batch_size
}
- /// Returns maximum number of rows in a row group.
+ /// Returns maximum number of rows in a row group, or `usize::MAX` if unlimited.
///
/// For more details see [`WriterPropertiesBuilder::set_max_row_group_size`]
pub fn max_row_group_size(&self) -> usize {
Review Comment:
That makes sense, as we also deprecate the corresponding setter. Done!
##########
parquet/src/file/properties.rs:
##########
@@ -45,6 +45,9 @@ pub const DEFAULT_STATISTICS_ENABLED: EnabledStatistics = EnabledStatistics::Pag
pub const DEFAULT_WRITE_PAGE_HEADER_STATISTICS: bool = false;
/// Default value for [`WriterProperties::max_row_group_size`]
pub const DEFAULT_MAX_ROW_GROUP_SIZE: usize = 1024 * 1024;
+/// Default value for [`WriterProperties::max_row_group_bytes`] (128 MB, same as the official Java
+/// implementation for `parquet.block.size`)
+pub const DEFAULT_MAX_ROW_GROUP_BYTES: usize = 128 * 1024 * 1024;
Review Comment:
> Or should we set it?
@alamb I think not, as it changes behaviour without users opting-in for that new behaviour. `None` preserves the existing behaviour by default, which is no byte count limit at all.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2770305033
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
+ RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
+ }
+
+ #[test]
+ fn test_row_group_limit_none_writes_single_row_group() {
+ // When both limits are None, all data should go into a single row group
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[1000],
+ "With no limits, all rows should be in a single row group"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_rows_only() {
+ // When only max_row_group_size is set, respect the row limit
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(300)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[300, 300, 300, 100],
+ "Row groups should be split by row count"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_bytes_only() {
+ // When only max_row_group_bytes is set, respect the byte limit
+ // Create batches with string data for more predictable byte sizes
+ // Write in multiple small batches so byte-based splitting can work
+ // (first batch establishes the avg row size, subsequent batches are split)
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "str",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+
+ // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(Some(3500))
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
+
+ // Write 10 batches of 10 rows each (100 rows total)
+ // Each string is ~100 bytes
+ for batch_idx in 0..10 {
+ let strings: Vec<String> = (0..10)
+ .map(|i| format!("{:0>100}", batch_idx * 10 + i))
+ .collect();
+ let array = StringArray::from(strings);
+ let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
+ writer.write(&batch).unwrap();
+ }
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ let sizes = row_group_sizes(builder.metadata());
+
+ assert!(
+ sizes.len() > 1,
+ "Should have multiple row groups due to byte limit, got {sizes:?}",
+ );
+
+ let total_rows: i64 = sizes.iter().sum();
+ assert_eq!(total_rows, 100, "Total rows should be preserved");
+ }
+
+ #[test]
+ fn test_row_group_limit_both_row_wins() {
Review Comment:
you don't need statistics, you can calculate it from the data types you need to encocde
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769462094
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
+ RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
+ }
+
+ #[test]
+ fn test_row_group_limit_none_writes_single_row_group() {
+ // When both limits are None, all data should go into a single row group
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[1000],
+ "With no limits, all rows should be in a single row group"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_rows_only() {
+ // When only max_row_group_size is set, respect the row limit
Review Comment:
the comment is not on the correct line
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
+ RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
+ }
+
+ #[test]
+ fn test_row_group_limit_none_writes_single_row_group() {
+ // When both limits are None, all data should go into a single row group
Review Comment:
the comment is not on the correct line
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769446078
##########
parquet/src/file/properties.rs:
##########
@@ -575,7 +595,34 @@ impl WriterPropertiesBuilder {
/// If the value is set to 0.
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
+ self.max_row_group_row_count = Some(value);
+ self
+ }
+
+ /// Sets maximum number of rows in a row group, or `None` for unlimited.
+ ///
+ /// # Panics
+ /// If the value is `Some(0)`.
Review Comment:
Please add comment on this and `set_max_row_group_bytes` and `set_max_row_group_size` what will be the behavior when both are set
##########
parquet/src/file/properties.rs:
##########
@@ -575,7 +595,34 @@ impl WriterPropertiesBuilder {
/// If the value is set to 0.
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
+ self.max_row_group_row_count = Some(value);
+ self
+ }
+
+ /// Sets maximum number of rows in a row group, or `None` for unlimited.
+ ///
+ /// # Panics
+ /// If the value is `Some(0)`.
+ pub fn set_max_row_group_row_count(mut self, value: Option<usize>) -> Self {
+ if let Some(v) = value {
+ assert!(v > 0, "Cannot have a 0 max row group size");
+ }
+ self.max_row_group_row_count = value;
+ self
+ }
+
+ /// Sets maximum size of a row group in bytes, or `None` for unlimited.
+ ///
+ /// Row groups are flushed when their estimated encoded size exceeds this threshold.
+ /// This is similar to parquet-mr's `parquet.block.size` behavior.
Review Comment:
parquet-mr is just the official java implementation for parquet, you can rewrite the comment to clarify that this match the official parquet Java implementation
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2770294974
##########
parquet/src/file/properties.rs:
##########
@@ -575,7 +595,34 @@ impl WriterPropertiesBuilder {
/// If the value is set to 0.
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
Review Comment:
Wait wait, no so fast, this is a breaking change, as clippy will fail for users, I was asking, it might be in a different pr, but open for discussion. if you keep it please update the pr description under changes to users
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2790331330
##########
parquet/src/file/properties.rs:
##########
@@ -575,7 +595,34 @@ impl WriterPropertiesBuilder {
/// If the value is set to 0.
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
Review Comment:
> Wait wait, no so fast, this is a breaking change, as clippy will fail for users, I was asking, it might be in a different pr, but open for discussion. if you keep it please update the pr description under changes to users
Yeah, in general I agree with @yonipeleg33 -- I don't think a clippy failure is a breaking change per-se -- the rust compiler will be happy to compile it. If downstream projects want to take a more strict "clippy must pass" stance I donthink that is technically an API breakage
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -331,18 +341,58 @@ impl<W: Write + Send> ArrowWriter<W> {
),
};
- // If would exceed max_row_group_size, split batch
- if in_progress.buffered_rows + batch.num_rows() > self.max_row_group_size {
- let to_write = self.max_row_group_size - 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);
+ if let Some(max_rows) = self.max_row_group_row_count {
+ if 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);
Review Comment:
Since this recurses, this could potentially blow out the stack with pathalogical inputs (e.g. a RecordBatch with 1M rows with a max_row_group_count of 1). I don't think it is necessary to fix now, I just wanted to point it out
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "etseidl (via GitHub)" <gi...@apache.org>.
etseidl commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2784218134
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -331,18 +341,57 @@ impl<W: Write + Send> ArrowWriter<W> {
),
};
- // If would exceed max_row_group_size, split batch
- if in_progress.buffered_rows + batch.num_rows() > self.max_row_group_size {
- let to_write = self.max_row_group_size - 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);
+ if let Some(max_rows) = self.max_row_group_row_count {
+ if 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);
+ }
+ }
+
+ // 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 {
+ if in_progress.buffered_rows > 0 {
+ let current_bytes = in_progress.get_estimated_total_bytes();
+
+ if current_bytes >= max_bytes {
+ self.flush()?;
+ return self.write(batch);
+ }
+
+ let avg_row_bytes = current_bytes / in_progress.buffered_rows;
+ if avg_row_bytes > 0 {
+ let remaining_bytes = max_bytes - current_bytes;
Review Comment:
Perhaps add a comment here that we've already checked that `current_bytes` is less than `max_bytes` by this line...at first glance I was thinking this should be `saturating_sub`.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "yonipeleg33 (via GitHub)" <gi...@apache.org>.
yonipeleg33 commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2791565131
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -331,18 +341,58 @@ impl<W: Write + Send> ArrowWriter<W> {
),
};
- // If would exceed max_row_group_size, split batch
- if in_progress.buffered_rows + batch.num_rows() > self.max_row_group_size {
- let to_write = self.max_row_group_size - 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);
+ if let Some(max_rows) = self.max_row_group_row_count {
+ if 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);
Review Comment:
Thanks for looking into it!
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -331,18 +341,58 @@ impl<W: Write + Send> ArrowWriter<W> {
),
};
- // If would exceed max_row_group_size, split batch
- if in_progress.buffered_rows + batch.num_rows() > self.max_row_group_size {
- let to_write = self.max_row_group_size - 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);
+ if let Some(max_rows) = self.max_row_group_row_count {
+ if 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);
Review Comment:
Thanks for looking into it!
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "yonipeleg33 (via GitHub)" <gi...@apache.org>.
yonipeleg33 commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769417566
##########
parquet/src/file/properties.rs:
##########
@@ -575,7 +595,34 @@ impl WriterPropertiesBuilder {
/// If the value is set to 0.
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
+ self.max_row_group_row_count = Some(value);
+ self
+ }
+
+ /// Sets maximum number of rows in a row group, or `None` for unlimited.
+ ///
+ /// # Panics
+ /// If the value is `Some(0)`.
+ pub fn set_max_row_group_row_count(mut self, value: Option<usize>) -> Self {
+ if let Some(v) = value {
+ assert!(v > 0, "Cannot have a 0 max row group size");
+ }
+ self.max_row_group_row_count = value;
+ self
+ }
+
+ /// Sets maximum size of a row group in bytes, or `None` for unlimited.
+ ///
+ /// Row groups are flushed when their estimated encoded size exceeds this threshold.
+ /// This is similar to parquet-mr's `parquet.block.size` behavior.
Review Comment:
Done
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769486009
##########
parquet/src/file/properties.rs:
##########
@@ -575,7 +595,34 @@ impl WriterPropertiesBuilder {
/// If the value is set to 0.
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
+ self.max_row_group_row_count = Some(value);
+ self
+ }
+
+ /// Sets maximum number of rows in a row group, or `None` for unlimited.
+ ///
+ /// # Panics
+ /// If the value is `Some(0)`.
+ pub fn set_max_row_group_row_count(mut self, value: Option<usize>) -> Self {
+ if let Some(v) = value {
+ assert!(v > 0, "Cannot have a 0 max row group size");
+ }
+ self.max_row_group_row_count = value;
+ self
+ }
+
+ /// Sets maximum size of a row group in bytes, or `None` for unlimited.
+ ///
+ /// Row groups are flushed when their estimated encoded size exceeds this threshold.
+ /// This is similar to the official `parquet.block.size` behavior.
+ ///
+ /// # Panics
+ /// If the value is `Some(0)`.
+ pub fn set_max_row_group_bytes(mut self, value: Option<usize>) -> Self {
+ if let Some(v) = value {
+ assert!(v > 0, "Cannot have a 0 max row group bytes");
+ }
Review Comment:
```suggestion
assert_ne!(value, Some(0), "Cannot have a 0 max row group bytes");
```
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769659979
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
+ RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
+ }
+
+ #[test]
+ fn test_row_group_limit_none_writes_single_row_group() {
+ // When both limits are None, all data should go into a single row group
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[1000],
+ "With no limits, all rows should be in a single row group"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_rows_only() {
+ // When only max_row_group_size is set, respect the row limit
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(300)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[300, 300, 300, 100],
+ "Row groups should be split by row count"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_bytes_only() {
+ // When only max_row_group_bytes is set, respect the byte limit
+ // Create batches with string data for more predictable byte sizes
+ // Write in multiple small batches so byte-based splitting can work
+ // (first batch establishes the avg row size, subsequent batches are split)
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "str",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+
+ // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(Some(3500))
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
+
+ // Write 10 batches of 10 rows each (100 rows total)
+ // Each string is ~100 bytes
+ for batch_idx in 0..10 {
+ let strings: Vec<String> = (0..10)
+ .map(|i| format!("{:0>100}", batch_idx * 10 + i))
+ .collect();
+ let array = StringArray::from(strings);
+ let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
+ writer.write(&batch).unwrap();
+ }
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ let sizes = row_group_sizes(builder.metadata());
+
+ assert!(
+ sizes.len() > 1,
+ "Should have multiple row groups due to byte limit, got {sizes:?}",
+ );
+
+ let total_rows: i64 = sizes.iter().sum();
+ assert_eq!(total_rows, 100, "Total rows should be preserved");
+ }
+
+ #[test]
+ fn test_row_group_limit_both_row_wins() {
+ // When both limits are set, the row limit triggers first
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(200) // Will trigger at 200 rows
+ .set_max_row_group_bytes(Some(1024 * 1024)) // 1MB - won't trigger for small int data
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[200, 200, 200, 200, 200],
+ "Row limit should trigger before byte limit"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_both_bytes_wins() {
+ // When both limits are set, the byte limit triggers first
+ // Write in multiple small batches so byte-based splitting can work
Review Comment:
According to the comment on the method, the way you write batches should not affect, only the config, that is if the byte based got hit first, it should write that, if the row hit first it should write that.
and also, it should work regardless of how you feed the data
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769648479
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
+ RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
+ }
+
+ #[test]
+ fn test_row_group_limit_none_writes_single_row_group() {
+ // When both limits are None, all data should go into a single row group
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[1000],
+ "With no limits, all rows should be in a single row group"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_rows_only() {
+ // When only max_row_group_size is set, respect the row limit
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(300)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[300, 300, 300, 100],
+ "Row groups should be split by row count"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_bytes_only() {
+ // When only max_row_group_bytes is set, respect the byte limit
+ // Create batches with string data for more predictable byte sizes
+ // Write in multiple small batches so byte-based splitting can work
+ // (first batch establishes the avg row size, subsequent batches are split)
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "str",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+
+ // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(Some(3500))
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
+
+ // Write 10 batches of 10 rows each (100 rows total)
+ // Each string is ~100 bytes
+ for batch_idx in 0..10 {
+ let strings: Vec<String> = (0..10)
+ .map(|i| format!("{:0>100}", batch_idx * 10 + i))
+ .collect();
+ let array = StringArray::from(strings);
+ let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
+ writer.write(&batch).unwrap();
+ }
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ let sizes = row_group_sizes(builder.metadata());
+
+ assert!(
+ sizes.len() > 1,
+ "Should have multiple row groups due to byte limit, got {sizes:?}",
+ );
+
+ let total_rows: i64 = sizes.iter().sum();
+ assert_eq!(total_rows, 100, "Total rows should be preserved");
+ }
+
+ #[test]
+ fn test_row_group_limit_both_row_wins() {
Review Comment:
Can you add a similar test when bytes win with the same structure as this, i.e. writing single large batch, but only changing the config (same test with only conf change)
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769467123
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
+ RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
+ }
+
+ #[test]
+ fn test_row_group_limit_none_writes_single_row_group() {
+ // When both limits are None, all data should go into a single row group
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[1000],
+ "With no limits, all rows should be in a single row group"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_rows_only() {
+ // When only max_row_group_size is set, respect the row limit
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(300)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[300, 300, 300, 100],
+ "Row groups should be split by row count"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_bytes_only() {
+ // When only max_row_group_bytes is set, respect the byte limit
+ // Create batches with string data for more predictable byte sizes
+ // Write in multiple small batches so byte-based splitting can work
+ // (first batch establishes the avg row size, subsequent batches are split)
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "str",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+
+ // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
Review Comment:
the comment is not on the correct line (it should be on the `Some(3500)` one)
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
+ RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
+ }
+
+ #[test]
+ fn test_row_group_limit_none_writes_single_row_group() {
+ // When both limits are None, all data should go into a single row group
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[1000],
+ "With no limits, all rows should be in a single row group"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_rows_only() {
+ // When only max_row_group_size is set, respect the row limit
+ let batch = create_test_batch(1000);
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(300)
+ .set_max_row_group_bytes(None)
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)).unwrap();
+
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(
+ &row_group_sizes(builder.metadata()),
+ &[300, 300, 300, 100],
+ "Row groups should be split by row count"
+ );
+ }
+
+ #[test]
+ fn test_row_group_limit_bytes_only() {
+ // When only max_row_group_bytes is set, respect the byte limit
+ // Create batches with string data for more predictable byte sizes
+ // Write in multiple small batches so byte-based splitting can work
+ // (first batch establishes the avg row size, subsequent batches are split)
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "str",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+
+ // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(None)
+ .set_max_row_group_bytes(Some(3500))
+ .build();
+
+ let file = tempfile::tempfile().unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
+
+ // Write 10 batches of 10 rows each (100 rows total)
+ // Each string is ~100 bytes
+ for batch_idx in 0..10 {
+ let strings: Vec<String> = (0..10)
+ .map(|i| format!("{:0>100}", batch_idx * 10 + i))
+ .collect();
+ let array = StringArray::from(strings);
+ let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
+ writer.write(&batch).unwrap();
+ }
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ let sizes = row_group_sizes(builder.metadata());
+
+ assert!(
+ sizes.len() > 1,
+ "Should have multiple row groups due to byte limit, got {sizes:?}",
+ );
+
+ let total_rows: i64 = sizes.iter().sum();
+ assert_eq!(total_rows, 100, "Total rows should be preserved");
+ }
+
+ #[test]
+ fn test_row_group_limit_both_row_wins() {
+ // When both limits are set, the row limit triggers first
Review Comment:
the comment is not on the correct line
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769458573
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4518,4 +4575,185 @@ mod tests {
assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+
+ /// Helper to create a test batch with the given number of rows.
+ /// Each row is approximately 4 bytes (one i32).
+ fn create_test_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "int",
+ ArrowDataType::Int32,
+ false,
+ )]));
+ let array = Int32Array::from((0..num_rows as i32).collect::<Vec<_>>());
Review Comment:
```suggestion
let array = Int32Array::from_iter(0..num_rows as i32);
```
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "yonipeleg33 (via GitHub)" <gi...@apache.org>.
yonipeleg33 commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2771029203
##########
parquet/src/file/properties.rs:
##########
@@ -575,7 +595,34 @@ impl WriterPropertiesBuilder {
/// If the value is set to 0.
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
Review Comment:
I do agree that this API should be deprecated. Thanks for pointing it out!
> this is a breaking change
This is inherently _not_ a breaking change; the purpose of marking APIs as deprecated is to warn users _before_ making a breaking change, without actually making this change.
This PR already calls for a minor bump due to the new APIs introduced; deprecating the old one does not change the version semantics for this PR.
> as clippy will fail for users
It's a [rustc](https://doc.rust-lang.org/reference/attributes/diagnostics.html#r-attributes.diagnostics.deprecated) warning:
> The deprecated attribute marks an item as deprecated. rustc will issue warnings on usage of #[deprecated] items
So unless users add `-D warnings`, compilation won't break.
> it might be in a different pr
I don't mind either way: Leaving it here or opening a new PR for deprecating the old API. LMK what's your preference and I'll do it.
> if you keep it please update the pr description under changes to users
Done!
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#issuecomment-3867066450
Thank you @yonipeleg33 -- sorry I forgot to submit my review from the other day when I reviewed this PR
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2769392681
##########
parquet/src/file/properties.rs:
##########
@@ -575,7 +595,34 @@ impl WriterPropertiesBuilder {
/// If the value is set to 0.
pub fn set_max_row_group_size(mut self, value: usize) -> Self {
assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
+ self.max_row_group_row_count = Some(value);
+ self
+ }
+
+ /// Sets maximum number of rows in a row group, or `None` for unlimited.
+ ///
+ /// # Panics
+ /// If the value is `Some(0)`.
+ pub fn set_max_row_group_row_count(mut self, value: Option<usize>) -> Self {
+ if let Some(v) = value {
+ assert!(v > 0, "Cannot have a 0 max row group size");
+ }
+ self.max_row_group_row_count = value;
+ self
+ }
+
+ /// Sets maximum size of a row group in bytes, or `None` for unlimited.
+ ///
+ /// Row groups are flushed when their estimated encoded size exceeds this threshold.
+ /// This is similar to parquet-mr's `parquet.block.size` behavior.
Review Comment:
parquet-mr is just the official java implementation for parquet, you can rewrite the comment to clarify that
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add max_row_group_bytes option to WriterProperties [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9357:
URL: https://github.com/apache/arrow-rs/pull/9357#discussion_r2790456602
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -331,18 +341,58 @@ impl<W: Write + Send> ArrowWriter<W> {
),
};
- // If would exceed max_row_group_size, split batch
- if in_progress.buffered_rows + batch.num_rows() > self.max_row_group_size {
- let to_write = self.max_row_group_size - 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);
+ if let Some(max_rows) = self.max_row_group_row_count {
+ if 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);
Review Comment:
I filed a ticket to track
- https://github.com/apache/arrow-rs/issues/9386
Given the prior code path for `max_row_group_size` also uses recursion I don't think this is a new bug introduced by this PR (though the max bytes path is now also susceptible to the same issue
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org