Improve to_arrow_batch_reader performance + use to_arrow_batch_reader in upsert to lower memory pressure

[koenvo]

Summary

This PR updates the upsert logic to use batch processing. The main goal is to prevent out-of-memory (OOM) issues when updating large tables by avoiding loading all data at once.

Note: This has only been tested against the unit tests—no real-world datasets have been evaluated yet.

This PR partially depends on functionality introduced in #1817.

---

Notes

• Duplicate detection across multiple batches is not possible with this approach.
• All data is read sequentially, which may be slower than the parallel read used by to_arrow. fixed using concurrent_tasks parameter

---

Performance Comparison

In setups with many small files, network and metadata overhead become the dominant factor. This impacts batch reading performance, as each file contributes relatively more overhead than payload. In the test setup used here, metadata access was the largest cost.

Using to_arrow_batch_reader (sequential):

• Scan: 9993.50 ms
• To list: 19811.09 ms

Using to_arrow (parallel):

• Scan: 10607.88 ms

[jayceslesar]

fwiw I think we should try to get this merged in at some point. Some ideas:

1. Make it a flag to use the batchreader or not, some users might have basically infinite memory
2. Is there a more optimal way to batch data? Thinking along the lines of using partitions although that may already happen under the hood

[koenvo]

fwiw I think we should try to get this merged in at some point. Some ideas:

1. Make it a flag to use the batchreader or not, some users might have basically infinite memory
2. Is there a more optimal way to batch data? Thinking along the lines of using partitions although that may already happen under the hood

I've been thinking about what I (as a developer) want. The answer is: set max memory usage.

Some ideas:

1. Determine which partitions can fit together in memory and batch load those together
2. Fetching of parquet files can happen parallel and only do loading sequential
3. Combine 1 and 2

[koenvo]

Did an update and ran a quick benchmark with different concurrent_tasks settings on to_arrow_batch_reader():

table = catalog.get_table("some_table")

# Benchmark loop
p = table.scan().to_arrow_batch_reader(concurrent_tasks=100)
for batch in tqdm.tqdm(p):
    print(pool.max_memory())

Results (including pool.max_memory()):

• concurrent_tasks=1 → 52it [00:06, 7.73it/s] | Max memory: 7.4 MB
• concurrent_tasks=10 → 391it [00:06, 61.98it/s] | Max memory: 36.3 MB
• concurrent_tasks=20 → 1412it [00:15, 83.54it/s] | Max memory: 147 MB
• concurrent_tasks=100 → 1030it [00:09, 106.84it/s] | Max memory: 1.76 GB

Some more testing (on 100mbit connection):

scan.to_arrow_batch_reader(concurrent_tasks=10)
2025-06-03 11:02:48.986 INFO Starting
2025-06-03 11:05:10.927 INFO Rows: 13584102
2025-06-03 11:05:10.927 INFO Memory usage: 78.4MB

scan.to_arrow()
2025-06-03 11:05:47.211 INFO Starting
2025-06-03 11:08:09.907 INFO Rows: 13584102
2025-06-03 11:08:09.907 INFO Memory usage: 11GB

Note: Performance also depends on the network connection.

[corleyma]

Rather than create your own threadpool executor here, I think you should use the ExecutorFactory defined elsewhere in the repo. It has a get_or_create method that prevents creating a new threadpool on every call, among other things.

[koenvo]

Ah thanks! Had this changed but forgot to push. Only need to make sure I get a pool with the correct max_workers set. Can't just use the regular get_or_create as that might have an incorrect number of workers.

[jayceslesar]

Does this preserve the ordering still? It looks like it did in to_table

[koenvo]

Yes. Executor.map maintains ordering by default. It first submits all jobs, and then waits for the result in original order.

[jayceslesar]

🤩

[koenvo]

Did another update to get rid of the concurrent_tasks argument. It now defaults to the max-workers Config.

I also refactored to_arrow to use to_arrow_batch_reader under the hood to prevent duplicate implementations of the same functionality.