Improve upsert 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.

[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.

[Fokko]

Make it a flag to use the batchreader or not, some users might have basically infinite memory

Typically, I'm not a big fan of this kind of flag, since it will delegate the responsibility to the user, and not every user might know the implications. Instead, let's optimize for the best experience out of the box.

[Fokko]

@koenvo Looks like there are some issues with the list vs large_list, do you have time to dig into the issue? In the original implementation we had cast operations to ensure that we keep the original types.

[koenvo]

@koenvo Looks like there are some issues with the list vs large_list, do you have time to dig into the issue? In the original implementation we had cast operations to ensure that we keep the original types.

I was kind of scared by the

        if property_as_bool(self._io.properties, PYARROW_USE_LARGE_TYPES_ON_READ, False):
            deprecation_message(
                deprecated_in="0.10.0",
                removed_in="0.11.0",
                help_message=f"Property `{PYARROW_USE_LARGE_TYPES_ON_READ}` will be removed.",
            )
            result = result.cast(arrow_schema)

Currently, I do the cast in my own application code, as I ran into the same issue but I though it was just my implementation. Had to add the cast.

And happy to look into this.

[Fokko]

Two minor comments, this looks great 👍

[Fokko]

Moving this forward, this looks great @koenvo 🙌