fix: lazily initialize BatchCoalescer in CoalescedShuffleReaderStream to avoid schema type mismatch

[phillipleblanc]

Summary

The BatchCoalescer inside CoalescedShuffleReaderStream was eagerly initialized with the declared schema from the execution plan. However, the actual IPC shuffle data may have different Arrow types — for example, string columns declared as LargeUtf8 in the plan schema but written as Utf8 by the CSV reader. When InProgressPrimitiveArray<T>::copy_rows() tries to downcast the source array, the type mismatch causes a panic:

Internal("\"primitive array\"")

Root Cause

In distributed (Ballista) query execution, shuffle data is written as Arrow IPC by Stage N and read by Stage N+1 via ShuffleReaderExec. The CoalescedShuffleReaderStream wraps the IPC reader output and coalesces small batches. It eagerly creates a LimitedBatchCoalescer using input.schema() — the declared schema from the plan. But the actual Arrow arrays in the IPC files may use different concrete types than what the plan declares, causing the BatchCoalescer to create InProgressPrimitiveArray<T> with the wrong T.

Fix

Lazily initialize the BatchCoalescer from the first actual batch's schema instead of the declared schema. This is the same pattern already applied to RepartitionExec in spiceai/datafusion#135.

[Copilot]

Pull request overview

This PR fixes a panic in Ballista’s shuffle reader path by avoiding eager LimitedBatchCoalescer initialization with the execution plan’s declared schema when the actual Arrow IPC shuffle batches contain differing concrete Arrow types.

Changes:

• Make CoalescedShuffleReaderStream lazily initialize its LimitedBatchCoalescer from the first non-empty incoming RecordBatch schema.
• Store batch_size / limit on the stream to support lazy coalescer creation and adjust polling logic to handle Option<LimitedBatchCoalescer>.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.