Skip to content

Improve upsert memory pressure #1995

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Fokko merged 20 commits into from
Jun 20, 2025
Merged

Improve upsert memory pressure #1995

Changes from 1 commit
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
7abfee9
Move actual implementation of upsert from Table to Transaction
koenvo Mar 19, 2025
db334ae
Fix some incorrect usage of schema
koenvo Mar 19, 2025
cebfda3
Write a test for upsert transaction
koenvo Mar 25, 2025
52fd35e
Add failing test for multiple upserts in same transaction
koenvo Mar 27, 2025
f336c0b
Fix test
koenvo Apr 2, 2025
07890ac
Add failing test
koenvo May 13, 2025
ae0e60f
Use Transaction.table_metadata when doing the data scan in upsert
koenvo May 13, 2025
ce8d9ef
Remove as it's resolved
koenvo May 13, 2025
5bdb0b8
Use to_arrow_batch_reader instead of to_arrow in upsert to prevent OO…
koenvo May 13, 2025
65fe36d
Filter rows to insert on each iteration instead of keeping a list of …
koenvo May 13, 2025
4614543
Merge branch 'main' into feat/use-batchreader-in-upsert
koenvo May 13, 2025
88a4ad2
Accept concurrent_tasks when fetching record_batches
koenvo Jun 2, 2025
f8acdb0
Use ExecutorFactory
koenvo Jun 3, 2025
119d92f
Simplify to_arrow to use the optimized to_record_batches
koenvo Jun 3, 2025
5d3a6aa
Fix for shutdown pool after doing a map
koenvo Jun 3, 2025
445845d
minor
koenvo Jun 3, 2025
6b8dace
Improve comments in to_record_batches
koenvo Jun 3, 2025
2cd2137
Make sure the 'infer the types when reading (#1669)' works again
koenvo Jun 14, 2025
8e6f5e9
Update pyiceberg/io/pyarrow.py
koenvo Jun 20, 2025
6813061
Update pyiceberg/io/pyarrow.py
koenvo Jun 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Fix for shutdown pool after doing a map
  • Loading branch information
@koenvo
koenvo committed Jun 3, 2025
commit 5d3a6aa012fb1f24988ea3756d875655075b3155
41 changes: 17 additions & 24 deletions pyiceberg/io/pyarrow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,6 @@

from __future__ import annotations

import concurrent.futures
import fnmatch
import functools
import itertools
Expand All @@ -36,7 +35,6 @@
import uuid
import warnings
from abc import ABC, abstractmethod
from concurrent.futures import Future
from copy import copy
from dataclasses import dataclass
from enum import Enum
Expand Down Expand Up @@ -71,7 +69,6 @@
FileType,
FSSpecHandler,
)
from sortedcontainers import SortedList

from pyiceberg.conversions import to_bytes
from pyiceberg.exceptions import ResolveError
Expand Down Expand Up @@ -1570,7 +1567,6 @@ def to_table(self, tasks: Iterable[FileScanTask]) -> pa.Table:
ResolveError: When a required field cannot be found in the file
ValueError: When a field type in the file cannot be projected to the schema type
"""

arrow_schema = schema_to_pyarrow(self._projected_schema, include_field_ids=False)

batches = self.to_record_batches(tasks)
Expand All @@ -1592,9 +1588,7 @@ def to_table(self, tasks: Iterable[FileScanTask]) -> pa.Table:

return result

def to_record_batches(
self, tasks: Iterable[FileScanTask]
) -> Iterator[pa.RecordBatch]:
def to_record_batches(self, tasks: Iterable[FileScanTask]) -> Iterator[pa.RecordBatch]:
"""Scan the Iceberg table and return an Iterator[pa.RecordBatch].

Returns an Iterator of pa.RecordBatch with data from the Iceberg table
Expand All @@ -1617,26 +1611,25 @@ def to_record_batches(
total_row_count = 0
executor = ExecutorFactory.get_or_create()

with executor as pool:
should_stop = False
for batches in pool.map(
lambda task: list(self._record_batches_from_scan_tasks_and_deletes([task], deletes_per_file)), tasks
):
for batch in batches:
current_batch_size = len(batch)
if self._limit is not None:
if total_row_count + current_batch_size >= self._limit:
yield batch.slice(0, self._limit - total_row_count)
limit_reached = False
for batches in executor.map(
lambda task: list(self._record_batches_from_scan_tasks_and_deletes([task], deletes_per_file)), tasks
):
for batch in batches:
current_batch_size = len(batch)
if self._limit is not None:
if total_row_count + current_batch_size >= self._limit:
yield batch.slice(0, self._limit - total_row_count)

# This break will also cancel all tasks in the Pool
should_stop = True
break
# This break will also cancel all tasks in the Pool
limit_reached = True
break

yield batch
total_row_count += current_batch_size
yield batch
total_row_count += current_batch_size

if should_stop:
break
if limit_reached:
break

def _record_batches_from_scan_tasks_and_deletes(
self, tasks: Iterable[FileScanTask], deletes_per_file: Dict[str, List[ChunkedArray]]
Expand Down