This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11349-9d811590898b00993654275febd58e1e3b09776a in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit c60c6ab7d4d4cb90811789945d3dc1161e73c78d Author: Bryan Dang <[email protected]> AuthorDate: Wed Oct 7 09:54:14 2026 +0000 feat(parquet): Support per-column data page row count limit (#11339) (#11349) ### Which issue does this PR fix? Closes #11339 ### Rationale & Proposed Changes `WriterProperties` currently supports configuring data page size limits both globally and per-column via `set_data_page_size_limit` and `set_column_data_page_size_limit`. However, `set_data_page_row_count_limit` could only be configured globally. This PR adds per-column support for the data page row count limit: 1. Adds `WriterPropertiesBuilder::set_column_data_page_row_count_limit(self, col: ColumnPath, value: usize) -> Self`. 2. Adds `WriterProperties::column_data_page_row_count_limit(&self, col: &ColumnPath) -> usize`. 3. Stores `data_page_row_count_limit` in `ResolvedColumnProperties` and uses it inside `ColumnWriterImpl::should_add_data_page()` and `base_batch_size`. 4. Adds unit tests in `properties.rs` and `column/writer/mod.rs`. ### How Has This Been Tested? - Added unit tests: - `test_writer_properties_column_data_page_row_count_limit` - `test_writer_properties_panic_on_zero_column_data_page_row_count_limit` - `test_column_writer_column_data_page_row_count_limit` - Ran full test suite: `cargo test -p parquet` (All unit tests & doctests passed). --- parquet/src/column/writer/mod.rs | 25 +++++++++++- parquet/src/file/properties.rs | 84 ++++++++++++++++++++++++++++++++++++---- 2 files changed, 100 insertions(+), 9 deletions(-) diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index d7268422c4..432c97bfe8 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -634,7 +634,7 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { // materialized slice to split and the per-mini-batch work is O(1), so we // can safely use a much larger batch size. let base_batch_size = if both_levels_compact && has_levels { - self.props.data_page_row_count_limit() + self.column_props.data_page_row_count_limit } else { self.props.write_batch_size() }; @@ -1209,7 +1209,7 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { return false; } - self.page_metrics.num_buffered_rows as usize >= self.props.data_page_row_count_limit() + self.page_metrics.num_buffered_rows as usize >= self.column_props.data_page_row_count_limit || self .encoder .estimated_data_page_size() @@ -3153,6 +3153,27 @@ mod tests { assert_eq!(other_values, vec![10]); } + #[test] + fn test_column_writer_column_data_page_row_count_limit() { + let props = Arc::new( + WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_data_page_row_count_limit(100) + .set_column_data_page_row_count_limit(ColumnPath::from("col"), 3) + .set_write_batch_size(1) + .build(), + ); + let data = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; + + let col_values = + write_and_collect_page_values(ColumnPath::from("col"), Arc::clone(&props), data); + let other_values = write_and_collect_page_values(ColumnPath::from("other"), props, data); + + assert_eq!(col_values, vec![3, 3, 3, 1]); + assert_eq!(other_values, vec![10]); + } + #[test] fn test_column_writer_caps_page_size_for_large_byte_array_values() { // Regression: the post-write data page byte limit check only fires diff --git a/parquet/src/file/properties.rs b/parquet/src/file/properties.rs index 435e542085..6fc7b822ad 100644 --- a/parquet/src/file/properties.rs +++ b/parquet/src/file/properties.rs @@ -244,7 +244,6 @@ enum OffsetIndexSetting { /// ``` #[derive(Debug, Clone)] pub struct WriterProperties { - data_page_row_count_limit: usize, write_batch_size: usize, max_row_group_row_count: Option<usize>, max_row_group_bytes: Option<usize>, @@ -338,7 +337,21 @@ impl WriterProperties { /// /// For more details see [`WriterPropertiesBuilder::set_data_page_row_count_limit`] pub fn data_page_row_count_limit(&self) -> usize { - self.data_page_row_count_limit + self.default_column_properties + .data_page_row_count_limit() + .unwrap_or(DEFAULT_DATA_PAGE_ROW_COUNT_LIMIT) + } + + /// Returns data page row count limit for a specific column. + /// + /// Takes precedence over [`Self::data_page_row_count_limit`]. + /// + /// Note: this is a best effort limit based on the write batch size. + pub fn column_data_page_row_count_limit(&self, col: &ColumnPath) -> usize { + resolve_data_page_row_count_limit( + self.column_override(col), + &self.default_column_properties, + ) } /// Returns configured batch size for writes. @@ -599,6 +612,7 @@ impl WriterProperties { statistics_enabled: resolve_statistics_enabled(column, default), write_page_header_statistics: resolve_write_page_header_statistics(column, default), data_page_size_limit: resolve_data_page_size_limit(column, default), + data_page_row_count_limit: resolve_data_page_row_count_limit(column, default), dictionary_page_size_limit: resolve_dictionary_page_size_limit(column, default), data_page_v2_compression_ratio_threshold: resolve_data_page_v2_compression_ratio_threshold(column, default), @@ -623,7 +637,6 @@ impl WriterProperties { /// See example on [`WriterProperties`] #[derive(Debug, Clone)] pub struct WriterPropertiesBuilder { - data_page_row_count_limit: usize, write_batch_size: usize, max_row_group_row_count: Option<usize>, max_row_group_bytes: Option<usize>, @@ -650,7 +663,6 @@ impl Default for WriterPropertiesBuilder { /// Returns default state of the builder. fn default() -> Self { Self { - data_page_row_count_limit: DEFAULT_DATA_PAGE_ROW_COUNT_LIMIT, write_batch_size: DEFAULT_WRITE_BATCH_SIZE, max_row_group_row_count: Some(DEFAULT_MAX_ROW_GROUP_ROW_COUNT), max_row_group_bytes: None, @@ -708,7 +720,6 @@ impl WriterPropertiesBuilder { } WriterProperties { - data_page_row_count_limit: self.data_page_row_count_limit, write_batch_size: self.write_batch_size, max_row_group_row_count: self.max_row_group_row_count, max_row_group_bytes: self.max_row_group_bytes, @@ -762,7 +773,8 @@ impl WriterPropertiesBuilder { /// If the value is `0`. pub fn set_data_page_row_count_limit(mut self, value: usize) -> Self { assert_ne!(value, 0, "Cannot have a 0 data page row count limit"); - self.data_page_row_count_limit = value; + self.default_column_properties + .set_data_page_row_count_limit(value); self } @@ -1323,6 +1335,18 @@ impl WriterPropertiesBuilder { self } + /// Sets data page row count limit for a specific column. + /// + /// Takes precedence over [`Self::set_data_page_row_count_limit`]. + /// + /// # Panics + /// If the value is `0`. + pub fn set_column_data_page_row_count_limit(mut self, col: ColumnPath, value: usize) -> Self { + assert_ne!(value, 0, "Cannot have a 0 data page row count limit"); + self.get_mut_props(col).set_data_page_row_count_limit(value); + self + } + /// Sets [`EnabledStatistics`] level for a specific column. /// /// Takes precedence over [`Self::set_statistics_enabled`]. @@ -1441,7 +1465,6 @@ impl WriterPropertiesBuilder { impl From<WriterProperties> for WriterPropertiesBuilder { fn from(props: WriterProperties) -> Self { WriterPropertiesBuilder { - data_page_row_count_limit: props.data_page_row_count_limit, write_batch_size: props.write_batch_size, max_row_group_row_count: props.max_row_group_row_count, max_row_group_bytes: props.max_row_group_bytes, @@ -1704,6 +1727,7 @@ struct ColumnProperties { encoding: Option<Encoding>, codec: Option<Compression>, data_page_size_limit: Option<usize>, + data_page_row_count_limit: Option<usize>, dictionary_page_size_limit: Option<usize>, dictionary_enabled: Option<bool>, statistics_enabled: Option<EnabledStatistics>, @@ -1743,6 +1767,12 @@ impl ColumnProperties { self.data_page_size_limit = Some(value); } + /// Sets data page row count limit for this column. + fn set_data_page_row_count_limit(&mut self, value: usize) { + assert_ne!(value, 0, "Cannot have a 0 data page row count limit"); + self.data_page_row_count_limit = Some(value); + } + /// Sets whether dictionary encoding is enabled for this column. fn set_dictionary_enabled(&mut self, enabled: bool) { self.dictionary_enabled = Some(enabled); @@ -1852,6 +1882,11 @@ impl ColumnProperties { self.data_page_size_limit } + /// Returns optional data page row count limit for this column. + fn data_page_row_count_limit(&self) -> Option<usize> { + self.data_page_row_count_limit + } + /// Returns optional statistics level requested for this column. If result is `None`, /// then no setting has been provided. fn statistics_enabled(&self) -> Option<EnabledStatistics> { @@ -1909,6 +1944,8 @@ pub(crate) struct ResolvedColumnProperties { pub(crate) write_page_header_statistics: bool, /// See [`WriterProperties::column_data_page_size_limit`]. pub(crate) data_page_size_limit: usize, + /// See [`WriterProperties::column_data_page_row_count_limit`]. + pub(crate) data_page_row_count_limit: usize, /// See [`WriterProperties::column_dictionary_page_size_limit`]. pub(crate) dictionary_page_size_limit: usize, /// See [`WriterProperties::column_data_page_v2_compression_ratio_threshold`]. @@ -1982,6 +2019,14 @@ fn resolve_data_page_size_limit( .unwrap_or(DEFAULT_PAGE_SIZE) } +fn resolve_data_page_row_count_limit( + column: Option<&ColumnProperties>, + default: &ColumnProperties, +) -> usize { + column_or_default(column, default, ColumnProperties::data_page_row_count_limit) + .unwrap_or(DEFAULT_DATA_PAGE_ROW_COUNT_LIMIT) +} + fn resolve_dictionary_page_size_limit( column: Option<&ColumnProperties>, default: &ColumnProperties, @@ -2619,6 +2664,31 @@ mod tests { ); } + #[test] + fn test_writer_properties_column_data_page_row_count_limit() { + let props = WriterProperties::builder() + .set_data_page_row_count_limit(100) + .set_column_data_page_row_count_limit(ColumnPath::from("col"), 10) + .build(); + + assert_eq!(props.data_page_row_count_limit(), 100); + assert_eq!( + props.column_data_page_row_count_limit(&ColumnPath::from("col")), + 10 + ); + assert_eq!( + props.column_data_page_row_count_limit(&ColumnPath::from("other")), + 100 + ); + } + + #[test] + #[should_panic(expected = "Cannot have a 0 data page row count limit")] + fn test_writer_properties_panic_on_zero_column_data_page_row_count_limit() { + let _ = WriterProperties::builder() + .set_column_data_page_row_count_limit(ColumnPath::from("col"), 0); + } + #[test] fn test_reader_properties_default_settings() { let props = ReaderProperties::builder().build();
