BatchCoalescer return error containing which columns failed to concatenate
- Dominant language
- Rust
- Stars
- 3.6k
- Forks
- 1.3k
- Avg merge
- 2d 16h
- Merged PRs (30d)
- 168
Description
**Is your feature request related to a problem or challenge? Please describe what you are trying to do.**
We're writing some code that combines multiple `RecordBatch`s using the `BatchCoalescer`.
Before we concatenate, we try to compute the correct size dictionary key to use for some columns in the `RecordBatch`. For example, if two batches are using a u8 key, but the total cardinality of the column would cause the dict to overflow, we would switch the column to a dict keyed by u16 before concatenating.
To do this, we're inspecting the columns to compute the total cardinality. However, we're thinking it might be faster in a lot of cases to just optimistically try to concat the `RecordBatch`s, and if it fails, then compute the right-sized dict key for the column and retry.
Unfortunately, the `BatchCoalescer` returns an opaque error and we can't tell for which column the dictionary overflowed.
```rs
let schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)), true),
Field::new("b", DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)), true),
]));
let mut dict1_avals = vec![];
let mut dict1_akeys = vec![];
for i in 0..256 {
dict1_akeys.push(i as u8);
dict1_avals.push(format!("{i}"))
}
let rb1 = RecordBatch::try_new(schema.clone(), vec![
Arc::new(DictionaryArray::new(
UInt8Array::from_iter_values(dict1_akeys),
Arc::new(StringArray::from_iter_values(dict1_avals))
)),
Arc::new(DictionaryArray::new(
UInt8Array::from_iter_values(vec![0;256]),
Arc::new(StringArray::from_iter_values(vec!["b"]))
)),
]).unwrap();
let rb2 = RecordBatch::try_new(schema.clone(), vec![
Arc::new(DictionaryArray::new(
UInt8Array::from_iter_values(vec![0]),
Arc::new(StringArray::from_iter_values(vec!["a"])),
)),
Arc::new(DictionaryArray::new(
UInt8Array::from_iter_values(vec![0]),
Arc::new(StringArray::from_iter_values(vec!["b"]))
)),
]).unwrap();
let mut batcher = arrow::compute::BatchCoalescer::new(
schema.clone(), rb1.num_rows() + rb2.num_rows()
);
batcher.push_batch(rb1).unwrap();
batcher.push_batch(rb2).unwrap(); // panics with opaque `DictionaryKeyOverflowError`
batcher.finish_buffered_batch().unwrap();
let result = batcher.next_completed_batch().unwrap();
```
**Describe the solution you'd like**
It would be nice if the error contained context about which column failed to concatenate. In the example above, we'd like to know that it was column "a" that caused the `DictionaryKeyOverflowError`.
I think we could maybe augment the error returned from `InProgressArray::finish` here to include the column name:
https://github.com/apache/arrow-rs/blob/aa626e12de8bc0d0f56b5349239cae1be8d1a195/arrow-select/src/coalesce.rs#L488-L496
**Describe alternatives you've considered**
**Additional context**
Contributor guide
Research direction
Start in arrow-select/src/coalesce.rs at InProgressArray::finish and trace how BatchCoalescer propagates its errors. Reproduce the two-column example from the issue, then verify that a dictionary overflow identifies the failing column, such as "a", in the returned error.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- data-engineering
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 72/100