This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 08651bc3d6 feat(parquet): Support per-column data page row count limit
(#11339) (#11349)
08651bc3d6 is described below
commit 08651bc3d68ceefd900586489a098a7eca770ea6
Author: Bryan Dang <[email protected]>
AuthorDate: Wed Oct 7 10:25:38 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();
