adriangb opened a new issue, #10931:
URL: https://github.com/apache/arrow-rs/issues/10931
### 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 [email protected] [email protected] 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<i64> {
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.*
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
