Skip to content

Add deferred materialization for oversized coalesced takes in BatchCoalescer #9760

Description

@ClSlaid

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?"
#9758 (comment)

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<RecordBatch> 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 takes.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions