apache / apache/arrow-rs

`parquet`: `write_batch_size` is bypassed for columns with compact level data, contrary to its documentation

Open
#10,931 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.