Skip to content

Commit b7dbf56

Browse files
committed
fix: stream record batches lazily in to_arrow_batch_reader()
1 parent ebbc0ba commit b7dbf56

3 files changed

Lines changed: 198 additions & 2 deletions

File tree

‎pyiceberg/io/pyarrow.py‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1744,6 +1744,24 @@ def _read_all_delete_files(io: FileIO, tasks: Iterable[FileScanTask]) -> dict[st
17441744
return deletes_per_file
17451745

17461746

1747+
def _read_deletes_for_task(io: FileIO, task: FileScanTask) -> dict[str, list[ChunkedArray]]:
1748+
"""Read the delete files belonging to a single file scan task.
1749+
1750+
Unlike `_read_all_delete_files`, this reads nothing up front for tasks the
1751+
consumer may never reach, so streaming readers never pay for delete files
1752+
ahead of the batch being yielded.
1753+
"""
1754+
deletes_per_file: dict[str, list[ChunkedArray]] = {}
1755+
for delete_file in task.delete_files:
1756+
for file, arr in _read_deletes(io, delete_file).items():
1757+
if file in deletes_per_file:
1758+
deletes_per_file[file].append(arr)
1759+
else:
1760+
deletes_per_file[file] = [arr]
1761+
1762+
return deletes_per_file
1763+
1764+
17471765
class ArrowScan:
17481766
_table_metadata: TableMetadata
17491767
_io: FileIO
@@ -1875,6 +1893,42 @@ def batches_for_task(task: FileScanTask) -> list[pa.RecordBatch]:
18751893
# This break will also cancel all running tasks in the executor
18761894
break
18771895

1896+
def to_record_batches_lazy(self, tasks: Iterable[FileScanTask]) -> Iterator[pa.RecordBatch]:
1897+
"""Stream record batches one file scan task at a time, in the calling thread.
1898+
1899+
Unlike `to_record_batches`, this never fans work out to the executor and
1900+
never reads ahead of the consumer: each task's delete files are read only
1901+
when the consumer reaches that task, and each task's batches are yielded
1902+
directly instead of being collected into a per-task list first. Peak
1903+
memory therefore stays flat no matter how many files the scan covers.
1904+
1905+
This backs `to_arrow_batch_reader()`, which documents low-memory
1906+
streaming. Callers that want maximum throughput (`to_table()`,
1907+
`to_pandas()`) should keep using `to_record_batches`.
1908+
1909+
Args:
1910+
tasks: FileScanTasks representing the data files and delete files to read from.
1911+
1912+
Returns:
1913+
An Iterator of PyArrow RecordBatches, in task order.
1914+
Total number of rows will be capped if specified.
1915+
1916+
Raises:
1917+
ResolveError: When a required field cannot be found in the file
1918+
ValueError: When a field type in the file cannot be projected to the schema type
1919+
"""
1920+
total_row_count = 0
1921+
for task in tasks:
1922+
deletes_per_file = _read_deletes_for_task(self._io, task)
1923+
for batch in self._record_batches_from_scan_tasks_and_deletes([task], deletes_per_file):
1924+
current_batch_size = len(batch)
1925+
if self._limit is not None and total_row_count + current_batch_size >= self._limit:
1926+
yield batch.slice(0, self._limit - total_row_count)
1927+
return
1928+
else:
1929+
yield batch
1930+
total_row_count += current_batch_size
1931+
18781932
def _record_batches_from_scan_tasks_and_deletes(
18791933
self, tasks: Iterable[FileScanTask], deletes_per_file: dict[str, list[ChunkedArray]]
18801934
) -> Iterator[pa.RecordBatch]:

‎pyiceberg/table/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2381,7 +2381,7 @@ def _to_arrow_batch_reader_via_file_scan_tasks(
23812381
scan.case_sensitive,
23822382
scan.limit,
23832383
dictionary_columns=dictionary_columns,
2384-
).to_record_batches(tasks)
2384+
).to_record_batches_lazy(tasks)
23852385

23862386
if dictionary_columns:
23872387
# schema_to_pyarrow returns plain types, but ArrowScan yields dictionary-encoded

‎tests/io/test_pyarrow.py‎

Lines changed: 143 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@
6363
Or,
6464
)
6565
from pyiceberg.expressions.literals import literal
66-
from pyiceberg.io import S3_RETRY_STRATEGY_IMPL, InputStream, OutputStream, load_file_io
66+
from pyiceberg.io import S3_RETRY_STRATEGY_IMPL, FileIO, InputFile, InputStream, OutputFile, OutputStream, load_file_io
6767
from pyiceberg.io.pyarrow import (
6868
ICEBERG_SCHEMA,
6969
PYARROW_PARQUET_FIELD_ID_KEY,
@@ -5486,3 +5486,145 @@ def test_dictionary_columns_produces_dict_encoded_output(tmpdir: str) -> None:
54865486

54875487
# Values must be identical
54885488
assert result_plain.column("label").to_pylist() == result_dict.column("label").to_pylist()
5489+
5490+
5491+
class _CountingFileIO(FileIO):
5492+
"""FileIO decorator recording how many files are opened for reading."""
5493+
5494+
def __init__(self, inner: FileIO) -> None:
5495+
super().__init__()
5496+
self._inner = inner
5497+
self.files_opened_for_read = 0
5498+
5499+
def new_input(self, location: str) -> InputFile:
5500+
self.files_opened_for_read += 1
5501+
return self._inner.new_input(location)
5502+
5503+
def new_output(self, location: str) -> OutputFile:
5504+
return self._inner.new_output(location)
5505+
5506+
def delete(self, location: str | InputFile | OutputFile) -> None:
5507+
self._inner.delete(location)
5508+
5509+
5510+
class _MockScan:
5511+
"""Minimal scan stub for `_to_arrow_batch_reader_via_file_scan_tasks`."""
5512+
5513+
def __init__(self, table_metadata: TableMetadataV2, io: FileIO, limit: int | None = None) -> None:
5514+
self.table_metadata = table_metadata
5515+
self.io = io
5516+
self.row_filter = AlwaysTrue()
5517+
self.case_sensitive = True
5518+
self.limit = limit
5519+
5520+
5521+
def _write_lazy_reader_test_files(
5522+
tmpdir: str, num_files: int, rows_per_file: int
5523+
) -> tuple[Schema, TableMetadataV2, list[FileScanTask]]:
5524+
"""Write `num_files` parquet files with consecutive, non-overlapping ids."""
5525+
iceberg_schema = Schema(
5526+
NestedField(1, "id", IntegerType(), required=False),
5527+
)
5528+
arrow_schema = pa.schema([pa.field("id", pa.int32(), nullable=True, metadata={PYARROW_PARQUET_FIELD_ID_KEY: "1"})])
5529+
table_metadata = TableMetadataV2(
5530+
location=f"file://{tmpdir}",
5531+
last_column_id=1,
5532+
format_version=2,
5533+
schemas=[iceberg_schema],
5534+
partition_specs=[PartitionSpec()],
5535+
)
5536+
tasks = []
5537+
for file_idx in range(num_files):
5538+
start_id = file_idx * rows_per_file
5539+
arrow_table = pa.table(
5540+
[pa.array(range(start_id, start_id + rows_per_file), type=pa.int32())],
5541+
schema=arrow_schema,
5542+
)
5543+
data_file = _write_table_to_data_file(f"{tmpdir}/lazy_reader_{file_idx}.parquet", arrow_schema, arrow_table)
5544+
data_file.spec_id = 0
5545+
tasks.append(FileScanTask(data_file))
5546+
return iceberg_schema, table_metadata, tasks
5547+
5548+
5549+
def test_to_arrow_batch_reader_does_not_read_ahead(tmpdir: str) -> None:
5550+
"""Regression test for https://github.com/apache/iceberg-python/issues/2407.
5551+
5552+
`to_arrow_batch_reader()` documents low-memory streaming ("a RecordBatch is
5553+
read one at a time"), but the reader used to fan every file scan task out to
5554+
the executor and materialize each file's batches into a list, so taking a
5555+
single batch from the reader read every file in the scan.
5556+
"""
5557+
from pyiceberg.table import _to_arrow_batch_reader_via_file_scan_tasks
5558+
5559+
num_files = 4
5560+
iceberg_schema, table_metadata, tasks = _write_lazy_reader_test_files(tmpdir, num_files, rows_per_file=5000)
5561+
io = _CountingFileIO(PyArrowFileIO())
5562+
5563+
reader = _to_arrow_batch_reader_via_file_scan_tasks(_MockScan(table_metadata, io), iceberg_schema, tasks)
5564+
5565+
first_batch = next(reader)
5566+
assert first_batch.num_rows > 0
5567+
assert io.files_opened_for_read == 1, (
5568+
f"expected exactly 1 file opened after consuming the first batch, got {io.files_opened_for_read}"
5569+
)
5570+
5571+
5572+
def test_to_record_batches_lazy_matches_eager(tmpdir: str) -> None:
5573+
"""The lazy streaming path must return the same rows, in the same order, as the threaded path."""
5574+
num_files = 3
5575+
rows_per_file = 2500
5576+
total_rows = num_files * rows_per_file
5577+
iceberg_schema, table_metadata, tasks = _write_lazy_reader_test_files(tmpdir, num_files, rows_per_file)
5578+
expected_ids = list(range(total_rows))
5579+
5580+
for limit in (None, 0, 1, 100, total_rows, total_rows + 10):
5581+
eager_scan = ArrowScan(table_metadata, PyArrowFileIO(), iceberg_schema, AlwaysTrue(), True, limit)
5582+
lazy_scan = ArrowScan(table_metadata, PyArrowFileIO(), iceberg_schema, AlwaysTrue(), True, limit)
5583+
5584+
eager_ids = [row for batch in eager_scan.to_record_batches(tasks) for row in batch.column("id").to_pylist()]
5585+
lazy_ids = [row for batch in lazy_scan.to_record_batches_lazy(tasks) for row in batch.column("id").to_pylist()]
5586+
5587+
assert lazy_ids == eager_ids, f"limit={limit}: lazy path diverged from the threaded path"
5588+
assert lazy_ids == expected_ids[: len(lazy_ids)], f"limit={limit}: unexpected row contents or ordering"
5589+
if limit is None or limit >= total_rows:
5590+
assert len(lazy_ids) == total_rows
5591+
else:
5592+
assert len(lazy_ids) == limit
5593+
5594+
5595+
def test_to_record_batches_lazy_applies_positional_deletes(tmpdir: str) -> None:
5596+
"""The lazy path must apply per-task positional deletes exactly like the eager path."""
5597+
from pyiceberg.table import _to_arrow_batch_reader_via_file_scan_tasks
5598+
5599+
iceberg_schema = Schema(
5600+
NestedField(1, "id", IntegerType(), required=False),
5601+
)
5602+
arrow_schema = pa.schema([pa.field("id", pa.int32(), nullable=True, metadata={PYARROW_PARQUET_FIELD_ID_KEY: "1"})])
5603+
table_metadata = TableMetadataV2(
5604+
location=f"file://{tmpdir}",
5605+
last_column_id=1,
5606+
format_version=2,
5607+
schemas=[iceberg_schema],
5608+
partition_specs=[PartitionSpec()],
5609+
)
5610+
5611+
data_file = _write_table_to_data_file(
5612+
f"{tmpdir}/lazy_reader_deletes.parquet",
5613+
arrow_schema,
5614+
pa.table([pa.array([1, 2, 3, 4], type=pa.int32())], schema=arrow_schema),
5615+
)
5616+
data_file.spec_id = 0
5617+
5618+
# Positional delete of row position 2 (value 3)
5619+
deletes_path = f"{tmpdir}/lazy_reader_deletes_pos.parquet"
5620+
pq.write_table(pa.table({"file_path": [data_file.file_path], "pos": [2]}), deletes_path)
5621+
delete_file = DataFile.from_args(
5622+
content=DataFileContent.POSITION_DELETES, file_path=deletes_path, file_format=FileFormat.PARQUET
5623+
)
5624+
tasks = [FileScanTask(data_file, delete_files={delete_file})]
5625+
5626+
result = _to_arrow_batch_reader_via_file_scan_tasks(
5627+
_MockScan(table_metadata, PyArrowFileIO()), iceberg_schema, tasks
5628+
).read_all()
5629+
5630+
assert result.column("id").to_pylist() == [1, 2, 4]

0 commit comments

Comments
 (0)