apache / apache/arrow-rs

Add deferred materialization for oversized coalesced takes in BatchCoalescer

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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.