feat: Add ArrivalOrder to ArrowScan for bounded-memory concurrent reads

[sumedhsakdeo]

Summary

Addresses #3036 — ArrowScan.to_record_batches() uses executor.map + list() which eagerly materializes all record batches per file into memory, causing OOM on large tables.

This PR adds a new order parameter to to_arrow_batch_reader() with two implementations:

• TaskOrder (default) — preserves existing behavior: batches grouped by file in task submission order, each file fully materialized before proceeding to the next.
• ArrivalOrder — yields batches as they are produced across files without materializing entire files into memory. Accepts three sub-parameters:
  • concurrent_streams: int — number of files to read concurrently (default: 8). A per-scan ThreadPoolExecutor(max_workers=concurrent_streams) bounds concurrency.
  • batch_size: int | None — number of rows per batch passed to PyArrow's ds.Scanner (default: PyArrow's built-in 131,072).
  • max_buffered_batches: int — size of the bounded queue between producers and consumer (default: 16), providing backpressure to cap memory usage.

Problem

The current implementation materializes all batches from each file via list() inside executor.map, which runs up to min(32, cpu_count+4) files in parallel. For large files this means all batches from ~20 files are held in memory simultaneously before any are yielded to the consumer.

Solution

Before: OOM on large tables

batches = table.scan().to_arrow_batch_reader()

After: bounded memory, tunable parallelism

from pyiceberg.table import ArrivalOrder

batches = table.scan().to_arrow_batch_reader(
    order=ArrivalOrder(concurrent_streams=4, batch_size=10000),
)

Default behavior is unchanged — TaskOrder preserves the existing executor.map + list() path for backwards compatibility.

Architecture

When order=ArrivalOrder(...), batches flow through _bounded_concurrent_batches:

1. All file tasks are submitted to a per-scan ThreadPoolExecutor(max_workers=concurrent_streams)
2. Workers push batches into a bounded Queue(maxsize=max_buffered_batches) — when full, workers block (backpressure)
3. The consumer yields batches from the queue via blocking queue.get()
4. A sentinel value signals completion — no timeout-based polling
5. On early termination (consumer stops), a cancel event is set and the queue is drained until the sentinel to unblock all stuck workers
6. The executor context manager handles deterministic shutdown

Refactored to_record_batches into helpers: _prepare_tasks_and_deletes, _iter_batches_arrival, _iter_batches_materialized, _apply_limit.

Ordering semantics

Configuration	File ordering	Within-file ordering
TaskOrder() (default)	Batches grouped by file, in task submission order	Row order
ArrivalOrder(concurrent_streams=1)	Grouped by file, sequential	Row order
ArrivalOrder(concurrent_streams>1)	Interleaved (no grouping guarantee)	Row order within each file

PR Stack

Breakdown of this large PR into smaller PRs:

1. PR 0: batch_size forwarding
2. PR 1: TaskOrder/ArrivalOrder enum — stop materializing entire files
3. PR 2: concurrent_streams — bounded concurrent reads in arrival order
4. PR 3: benchmark + docs guidance

Benchmark results

32 files × 500K rows, 5 columns (int64, float64, string, bool, timestamp), batch_size=131,072 (PyArrow default):

Config	Throughput (rows/s)	TTFR (ms)	Peak Arrow Memory
default (TaskOrder)	190,250,192	73.4	642.2 MB
ArrivalOrder(cs=1)	59,317,085	27.7	10.3 MB
ArrivalOrder(cs=2)	105,414,909	28.8	42.0 MB
ArrivalOrder(cs=4)	175,840,782	28.4	105.5 MB
ArrivalOrder(cs=8)	211,922,538	32.3	271.7 MB
ArrivalOrder(cs=16)	209,011,424	45.0	473.3 MB

TTFR = Time to First Record, cs = concurrent_streams

Note on throughput plateau at cs=8: This benchmark runs against local filesystem where Parquet reads are CPU-bound (decompression + decoding). Throughput plateaus once enough threads saturate available cores. On cloud storage (S3/GCS/ADLS), reads are I/O-bound with 50-200ms per-file latency, so higher concurrent_streams values (16-64+) would continue to show throughput gains until network bandwidth saturates. The optimal concurrent_streams will be higher for remote storage than what this local benchmark suggests.

Positional deletes, row filters, and limit are handled correctly in all modes.

Are these changes tested?

Yes. 25 new unit tests across two test files, plus a micro-benchmark:

• tests/io/test_pyarrow.py (16 tests): batch_size controls rows per batch, arrival order yields all rows correctly, arrival order respects limit, within-file ordering preserved, positional deletes applied correctly in all three modes (task order, arrival order, concurrent), positional deletes with limit, concurrent_streams < 1 raises ValueError
• tests/io/test_bounded_concurrent_batches.py (10 tests): single/multi-file correctness, incremental streaming, backpressure blocks producers when queue is full, error propagation from workers to consumer, early termination cancels workers cleanly, no deadlock when concurrent_streams > max_buffered_batches on early termination, concurrency limit enforced, empty task list, ArrowScan integration with limit
• tests/benchmark/test_read_benchmark.py: read throughput micro-benchmark across 6 configurations measuring rows/sec, TTFR, and peak Arrow memory

Are there any user-facing changes?

Yes. New order parameter on DataScan.to_arrow_batch_reader():

• order: ScanOrder | None — controls batch ordering. Pass TaskOrder() (default) or ArrivalOrder(concurrent_streams=N, batch_size=B, max_buffered_batches=M).

New public classes TaskOrder and ArrivalOrder (subclasses of ScanOrder) exported from pyiceberg.table.

All parameters are optional with backwards-compatible defaults. Existing code is unaffected.

Documentation updated in mkdocs/docs/api.md with usage examples, ordering semantics, and configuration guidance table.

[ShreyeshArangath]

Overall LGTM, just a few comments

[kevinjqliu]

Thanks for the PR! This is touching a lot of different areas. I wonder if theres a simpler solution to resolve #3036. It seems like we just need to stop calling list() on iterators.

Let me take a closer look at the problem

[sumedhsakdeo]

Thanks @kevinjqliu! You're right that removing list() is the core fix — and ArrivalOrder does exactly that. Two reasons I kept TaskOrder as the default rather than just swapping list() out:

1. Semantic change: today's behavior groups batches by file in submission order. Replacing list() with interleaved streaming would be a silent breaking change for any code that relies on per-file grouping. TaskOrder preserves the existing contract; ArrivalOrder opts in.

2. No backpressure: executor.map submits all N tasks immediately even with lazy iteration — 1000 files still means 1000 concurrent threads. ArrivalOrder(concurrent_streams=N) bounds both concurrency and memory via a bounded queue.

Time-to-first-record (TTFR) is particularly critical for ML training workloads where the GPU stalls waiting for the first batch — executor.map blocks on file 0 completing before yielding anything, even if files 1–15 are already done (head-of-line blocking). On simulated S3 workloads (variable TTFB), ArrivalOrder shows 7x lower TTFR and ~11% higher throughput vs TaskOrder at the same thread count. Locally it's ~15% faster for the same reason.

The PR is broken into a stack if you'd prefer to review incrementally — happy to land them in order.

[geruh]

Thanks for raising @sumedhsakdeo! I took an initial pass on this and If I'm following this correctly there seems to be two parts here, the bounded memory fix, and the "arrival ordering" optimization. To me these seem like two separate concerns that could be broken up.

However, I'm not sure about ScanOrder abstraction from a user perspective. Right now this is fundamentally just a choice on whether or not to stream the tasks with a bounded queue. I took a look at a few other projects and noticed pyarrow uses params like fragment_readahead, and batch_readahead. From my understanding ArrivalOrder(max_buffered_batches=N, concurrent_streams=M) is pretty similar to the configs above, and TaskOrder is similar to what we get with the executor logic today. Which makes it seem like we can consider a simpler param approach, unless there is a future strategy that motivates us to keep the abstraction.

Could the ordered path also use bounded streaming? You'd block on a slow file, but users would get the memory fix without giving up deterministic ordering.

[geruh]

nit: this is a cosmetic change

[geruh]

I think you could just rebase to reduce the noise here

[geruh]

Can you mark with @pytest.mark.benchmark

[Fokko]

Sorry, but I'm not a bit fan of this change because it adds a lot of threading/locking.

I think it should be possible to do a full streaming of the tasks, where we open up the manifests sequentially.

[kevinjqliu]

I took some time to understand how ArrowScan is implemented and used today; and also where the potential bottlenecks are. Here's a summary:

ArrowScan has 2 public functions for reading; to_table and to_record_batches.

• to_table materializes the entire table as an arrow table
• to_record_batches return an iterator with the expectation of lazy materialization

Let's focus on to_record_batches since to_table can just be implemented by consuming all the record batches.
We're looking at this problem both in terms of throughput (how much data can be read) and memory usage (how much data is buffered in memory).

ArrowScan can utilize multiple threads, via the ExecutorFactory, to spread out tasks to worker threads.
I realized a few things while reviewing the original code:

1. There's an implicit ordering of the record batches; the ordering is deterministic based on task order (task 0, task 1, ...). Batches from later tasks are not emitted before earlier tasks complete.
2. Because of the strict ordering and with multiple worker threads, there's inherent head-of-line blocking (as mentioned above). Later tasks may finish work, but their results cannot be emitted until earlier tasks complete. This blocking diminishes the overall throughput of ArrowScan's ability to produce record batches.
3. Each worker materializes the entire task/file to memory (via list(self._record_batches_from_scan_tasks_and_deletes)) before returning results. This causes higher than expected memory utilization and imo is a bug because it's not the expected behavior for to_record_batches's lazy materialization.

To summarize, the current implementation fan out tasks to different workers. Each worker will read/materialize the entire task/file in memory before returning record batches. But only the first worker's record batches are consumed while the other workers wait.

The problem I would really like to solve immediately is (3), which drastically increases memory utilization (by # of workers * size of file in memory). I think the right behavior for to_record_batches is for each worker to not materialize the entire task/file in memory.

I think the bounded queue is a great idea for maximizing throughput while bounding memory usage. But it would only be useful when we relax the constraint for ordering. It would be great to brainstorm ways to introduce this change into the codebase while ensuring that it's maintainable.

[kevinjqliu]

My first instinct is to create a new Scan class which implements the bounded queue. This will isolate the changes to the new class and we just need to expose a way for table to execute reads with this implementation.
The new Scan class will allow us to be explicit about trading off strict ordering in order to maximize throughput (while also bounding memory usage)

[kevinjqliu]

created a discussion thread #3122