Add deferred materialization for oversized coalesced takes in BatchCoalescer
- Dominant language
- Rust
- Stars
- 3.6k
- Forks
- 1.3k
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 167
Description
**Is your feature request related to a problem or challenge? Please describe what you are trying to do.**
This is a follow-up to the review on `#9758`:
> @Dandandan: \"Shouldn't it cap `indices` at limit?\"
> https://github.com/apache/arrow-rs/pull/9758#discussion_r3104937679
`BatchCoalescer::push_batch_with_indices` now chunks oversized `indices`, which fixes the immediate issue. However, for fallback schemas such as `Utf8`, the coalescer still eagerly materializes all indexed `take` chunks up front.
This matters because `take` can amplify output size:
- the input batch may have `N` rows
- `indices.len()` may be much larger than `N`
- indices may repeat rows many times
For very large repeated takes, this fallback path is still noticeably expensive.
**Describe the solution you'd like**
Add deferred materialization for oversized indexed takes in `BatchCoalescer`.
High-level idea:
- when `push_batch_with_indices` falls back to materialized `take`
- enqueue the pending indexed-take work instead of evaluating all chunks immediately
- keep only a small ready window of completed batches
- materialize more only as completed batches are consumed
A reasonable first step may be to enable this only for fallback schemas such as `Utf8`, while leaving the existing direct path for primitive and view types unchanged.
**Describe alternatives you've considered**
- Keep the current eager chunked fallback
- Restrict lazy materialization to fallback types only
- Revisit API/error handling first, since deferred `take_record_batch(...)` work is awkward with `next_completed_batch()` returning `Option` instead of `Result<...>`
**Additional context**
We added an oversized-take benchmark with:
- input batch size: `8192`
- output indices length: `131072`
- `biggest_coalesce_batch_size = Some(1024)`
Current eager chunked behavior:
- `primitive extra_large_repeat`: about `1.23 ms`
- `mixed_utf8 extra_large_repeat`: about `76.6-81.2 ms`
Experimental lazy prototype in a separate worktree:
- `primitive extra_large_repeat`: about `1.32 ms`
- `mixed_utf8 extra_large_repeat`: about `29.5 ms`
These numbers are local and still imperfect, but they suggest deferred materialization could significantly improve oversized fallback `take`s.
Contributor guide
Research direction
Start at BatchCoalescer::push_batch_with_indices and run the oversized-take benchmark described in the issue, comparing primitive and mixed_utf8 cases. Trace next_completed_batch() and the fallback take path, then verify that oversized fallback work is materialized incrementally while the existing direct path remains unchanged and benchmark behavior improves.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- data-engineering
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100