apache / apache/arrow-rs

Parquet error: index overflow decoding byte array

Open
#7,243 2 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Rust
Stars
3.6k
Forks
1.3k
Avg merge
2d 18h
Merged PRs (30d)
169

Description

**Describe the bug**

When reading a dictionary-encoded string column from Parquet, we are getting an error from an index overflow reading data that was within valid bounds when written.

**To Reproduce**

See the below minimal reproduction:

```rust
use std::sync::Arc;

use arrow::{
array::{DictionaryArray, PrimitiveArray, RecordBatch, StringArray},
datatypes::{DataType, Field, Int32Type, Schema},
};
use bytes::{BufMut, BytesMut};
use parquet::arrow::{
ArrowWriter,
arrow_reader::{ArrowReaderMetadata, ParquetRecordBatchReaderBuilder},
};

fn main() {
// Schema is single column, dictionary-encoded string
let schema = Arc::new(Schema::new(vec![Field::new(
"a",
DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
true,
)]));

// First batch contains a single non-null row, with a massive 4MB string
let batch1 = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(DictionaryArray::new(
PrimitiveArray::::from(vec![Some(0)]),
Arc::new(StringArray::from(vec![Some("a".repeat(4_000_000))])),
))],
)
.unwrap();

// Second batch is 600 nulls
let batch2 = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(DictionaryArray::new(
PrimitiveArray::::from(vec![None; 600]),
Arc::new(StringArray::new_null(0)),
))],
)
.unwrap();

// Third batch is a single non-null row, it doesn't matter what value is for this
// reproduction.
let batch3 = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(DictionaryArray::new(
PrimitiveArray::::from(vec![Some(0)]),
Arc::new(StringArray::from(vec![Some("b")])),
))],
)
.unwrap();

// Write the batches to parquet format
let buf = BytesMut::new();
let mut writer = ArrowWriter::try_new(buf.writer(), schema.clone(), None).unwrap();
for batch in [batch1, batch2, batch3] {
// Flush after each batch, this has the effect of placing each batch into its own row group
writer.write(&batch).unwrap();
writer.flush().unwrap();
}
let buf = writer.into_inner().unwrap().into_inner().freeze();

// Sanity check the metadata
let meta = ArrowReaderMetadata::load(&buf, Default::default()).unwrap();
let row_group_meta = meta.metadata().row_groups();
let mut running_total = 0;
for (i, rg_meta) in row_group_meta.iter().enumerate() {
running_total += rg_meta.num_rows();
println!(
"row_group[{i}]={:>4}, total={running_total}",
rg_meta.num_rows()
);
}

// Read the parquet file back
let mut reader = ParquetRecordBatchReaderBuilder::try_new(buf)
.unwrap()
.build()
.unwrap();

// default batch size is large enough that the parquet reader merges the row groups into a
// single batch, and this merging op triggers the index overflow
reader.next().unwrap().unwrap();
}
```

**Expected behavior**

The reader should not overflow.

**Additional context**

I have done some debugging to arrive at this reproduction. Here's my analysis:

- The parquet reader will merge row groups to arrive at the desired batch size. In this case, 602 values is smaller than the default batch size, so it will read all three row groups as a single `RecordBatch`.
- Doing so requires merging the three dictionaries
- The second dictionary contains 600 nulls, so when reading the third row group the state is 601 values, of which only the first is non-null, a 4MB value.
- To merge the third dictionary, the current state is flattened. This has the effect of allocating 600*4MB chunks in the buffer for the nulls, which overflows the i32 size.

This is obviously a bit of an edge case, however we are hitting this in production at Pydantic.

We can do mitigations like truncate the 4MB value, limit batch sizes and also maybe switch from dictionaries to string views.

I think maybe there's also a possible mitigation in the `parquet` crate, to add a reader option to not merge row groups? That might be a nice efficiency win for decoding.

Contributor guide

Open the contributing guide

Research direction

Start with the parquet arrow reader entry points shown in the reproduction, especially ArrowReaderMetadata and ParquetRecordBatchReaderBuilder, then trace how row groups and dictionaries are merged into one RecordBatch. Re-run the supplied reproduction and inspect the dictionary flattening path. Done means the valid input reads without an index overflow and the regression is covered by a test.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
data-engineering
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.