`parquet`: `write_batch_size` is bypassed for columns with compact level data, contrary to its documentation
- Dominant language
- Rust
- Stars
- 3.6k
- Forks
- 1.3k
- Avg merge
- 2d 16h
- Merged PRs (30d)
- 168
Description
### Describe the bug
`WriterPropertiesBuilder::set_write_batch_size` documents `write_batch_size` as the granularity at which the writer's other limits are enforced (`parquet/src/file/properties.rs:735-743`, `main` @ `0aece99d08`):
> Sets write batch size (defaults to 1024 via [`DEFAULT_WRITE_BATCH_SIZE`]).
>
> For performance reasons, data for each column is written in batches of this size.
>
> Additional limits such as such as [`set_data_page_row_count_limit`] are checked between batches, and thus the write batch size value acts as an upper-bound on the enforcement granularity of other limits.
`set_data_page_row_count_limit` (`properties.rs:724-725`) and `set_data_page_size_limit` (`properties.rs:1113-1114`) both restate this as "Note: this is a best effort limit based on value of `set_write_batch_size`".
`GenericColumnWriter::write_batch_internal` does not read `write_batch_size` when the column's level data is compact (`parquet/src/column/writer/mod.rs:629-641`):
```rust
let both_levels_compact = !matches!(def_levels, LevelDataRef::Materialized(_))
&& !matches!(rep_levels, LevelDataRef::Materialized(_));
let has_levels = !matches!(def_levels, LevelDataRef::Absent)
|| !matches!(rep_levels, LevelDataRef::Absent);
// When both level vectors are compact (Uniform or Absent), there is no
// 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()
} else {
self.props.write_batch_size()
};
```
`LevelDataRef::Uniform` (`mod.rs:382-386`) is what the Arrow writer produces for a flat, non-repeated column whose definition levels are all the same value, which is the ordinary case of a nullable column whose batch happens to contain no nulls. Such a column takes the first branch, so its mini-batch size is `data_page_row_count_limit` (default 20,000) and `write_batch_size` is never consulted. `set_write_batch_size` then has no observable effect on that column.
Because the condition is on the level *data* and not on the schema, the same column with the same writer properties switches branches depending on whether a given batch happens to contain a null.
### To Reproduce
```
cargo new repro && cd repro
cargo add arrow@59.2.0 parquet@59.2.0 bytes@1
```
```rust
// src/main.rs
use arrow::array::{ArrayRef, Int64Array};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use parquet::arrow::ArrowWriter;
use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader};
use parquet::file::properties::WriterProperties;
use std::sync::Arc;
const ROWS: i64 = 12_000;
fn rows_per_page(nullable: bool, with_null: bool, write_batch_size: usize) -> Vec {
let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, nullable)]));
let col: Int64Array = (0..ROWS)
.map(|i| if with_null && i == 0 { None } else { Some(i) })
.collect();
let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(col) as ArrayRef]).unwrap();
let props = WriterProperties::builder()
.set_write_batch_size(write_batch_size)
.set_data_page_row_count_limit(3_000)
.set_dictionary_enabled(false)
.build();
let mut buf = Vec::new();
let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap();
writer.write(&batch).unwrap();
writer.close().unwrap();
let md = ParquetMetaDataReader::new()
.with_page_index_policy(PageIndexPolicy::Required)
.parse_and_finish(&bytes::Bytes::from(buf))
.unwrap();
let total = md.row_group(0).num_rows();
let pages = &md.offset_index().unwrap()[0][0].page_locations;
(0..pages.len())
.map(|i| pages.get(i + 1).map_or(total, |n| n.first_row_index) - pages[i].first_row_index)
.collect()
}
fn main() {
println!("{ROWS} Int64 rows in one batch, data_page_row_count_limit = 3000, no dictionary\n");
for (label, nullable, with_null) in [
("nullable=false (def levels Absent) ", false, false),
("nullable=true, no null in batch (def levels Uniform) ", true, false),
("nullable=true, one null in batch (def levels Materialized)", true, true),
] {
for wbs in [8usize, 64, 1024, 100_000] {
println!(
"{label} write_batch_size={wbs:<6} -> {:?}",
rows_per_page(nullable, with_null, wbs)
);
}
println!();
}
}
```
```
cargo run --release
```
Actual output on `parquet` 59.2.0:
```
12000 Int64 rows in one batch, data_page_row_count_limit = 3000, no dictionary
nullable=false (def levels Absent) write_batch_size=8 -> [3000, 3000, 3000, 3000]
nullable=false (def levels Absent) write_batch_size=64 -> [3008, 3008, 3008, 2976]
nullable=false (def levels Absent) write_batch_size=1024 -> [3072, 3072, 3072, 2784]
nullable=false (def levels Absent) write_batch_size=100000 -> [12000]
nullable=true, no null in batch (def levels Uniform) write_batch_size=8 -> [3000, 3000, 3000, 3000]
nullable=true, no null in batch (def levels Uniform) write_batch_size=64 -> [3000, 3000, 3000, 3000]
nullable=true, no null in batch (def levels Uniform) write_batch_size=1024 -> [3000, 3000, 3000, 3000]
nullable=true, no null in batch (def levels Uniform) write_batch_size=100000 -> [3000, 3000, 3000, 3000]
nullable=true, one null in batch (def levels Materialized) write_batch_size=8 -> [3000, 3000, 3000, 3000]
nullable=true, one null in batch (def levels Materialized) write_batch_size=64 -> [3008, 3008, 3008, 2976]
nullable=true, one null in batch (def levels Materialized) write_batch_size=1024 -> [3072, 3072, 3072, 2784]
nullable=true, one null in batch (def levels Materialized) write_batch_size=100000 -> [12000]
```
The first and third groups behave as documented: the row count limit is enforced only at mini-batch boundaries, so pages end at the first multiple of `write_batch_size` at or past 3,000, and a `write_batch_size` of 100,000 defeats the limit entirely. The middle group is unchanged across a 12,500x range of `write_batch_size`.
### Expected behavior
Either
- `write_batch_size` continues to bound the mini-batch size on this path, as documented (for example by taking the smaller of the two values), or
- the documentation on `set_write_batch_size`, `set_data_page_row_count_limit` and `set_data_page_size_limit` describes the real rule: that `write_batch_size` applies only to columns with materialized level data, and that a column with uniform or absent levels batches at `data_page_row_count_limit` instead.
I have no view on which is the right resolution; the two are inconsistent today and I could not tell which was intended from the code.
### Additional context
`parquet` 59.2.0 (crates.io), also present on `main` @ `0aece99d08`. The branch was introduced in #9831.
The substitution is not uniformly coarser than the documented behaviour. With the row count limit effectively disabled and `set_data_page_size_limit(10_000)` on the same Int64 column, 100,000 rows produce:
```
nullable=false (Absent) wbs=8 -> 80 pages, first page 1256 rows / 10070 bytes
nullable=false (Absent) wbs=1024 -> 49 pages, first page 2048 rows / 16406 bytes
nullable=true, no nulls (Uniform) wbs=8 -> 80 pages, first page 1250 rows / 10029 bytes
nullable=true, no nulls (Uniform) wbs=1024 -> 80 pages, first page 1250 rows / 10029 bytes
```
Here the compact path holds the byte limit more tightly than `write_batch_size=1024` does, because `ByteBudgetChunker` sizes the sub-batch from the page byte budget once `base_batch_size` is large enough to matter. The point is not that one branch is worse, but that the knob the documentation names as the enforcement granularity is not the one in effect.
---
*This issue was written by Claude (Anthropic's AI assistant) working with @adriangb. The reproduction above was executed and its output is verbatim.*
Contributor guide
Research direction
Reproduce the discrepancy with the supplied ArrowWriter program, then read parquet/src/file/properties.rs:724-743 and 1113-1114 alongside parquet/src/column/writer/mod.rs:629-641 and the LevelDataRef definitions at mod.rs:382-386. Determine which batching rule is intended for Uniform and Absent levels; done means the implementation and the documented limits agree, with coverage for the compact-level case.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- data
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100