fix: stream record batches lazily in to_arrow_batch_reader()

Author: gsrujankumar

pyiceberg/io/pyarrow.py

@@ -1760,6 +1760,24 @@ def _read_all_delete_files(io: FileIO, tasks: Iterable[FileScanTask]) -> dict[st
     return deletes_per_file
 
 
+def _read_deletes_for_task(io: FileIO, task: FileScanTask) -> dict[str, list[ChunkedArray]]:
+    """Read the delete files belonging to a single file scan task.
+
+    Unlike `_read_all_delete_files`, this reads nothing up front for tasks the
+    consumer may never reach, so streaming readers never pay for delete files
+    ahead of the batch being yielded.
+    """
+    deletes_per_file: dict[str, list[ChunkedArray]] = {}
+    for delete_file in task.delete_files:
+        for file, arr in _read_deletes(io, delete_file).items():
+            if file in deletes_per_file:
+                deletes_per_file[file].append(arr)
+            else:
+                deletes_per_file[file] = [arr]
+
+    return deletes_per_file
+
+
 class ArrowScan:
     _table_metadata: TableMetadata
     _io: FileIO
@@ -1891,6 +1909,42 @@ def batches_for_task(task: FileScanTask) -> list[pa.RecordBatch]:
                 # This break will also cancel all running tasks in the executor
                 break
 
+    def to_record_batches_lazy(self, tasks: Iterable[FileScanTask]) -> Iterator[pa.RecordBatch]:
+        """Stream record batches one file scan task at a time, in the calling thread.
+
+        Unlike `to_record_batches`, this never fans work out to the executor and
+        never reads ahead of the consumer: each task's delete files are read only
+        when the consumer reaches that task, and each task's batches are yielded
+        directly instead of being collected into a per-task list first. Peak
+        memory therefore stays flat no matter how many files the scan covers.
+
+        This backs `to_arrow_batch_reader()`, which documents low-memory
+        streaming. Callers that want maximum throughput (`to_table()`,
+        `to_pandas()`) should keep using `to_record_batches`.
+
+        Args:
+            tasks: FileScanTasks representing the data files and delete files to read from.
+
+        Returns:
+            An Iterator of PyArrow RecordBatches, in task order.
+            Total number of rows will be capped if specified.
+
+        Raises:
+            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
+        """
+        total_row_count = 0
+        for task in tasks:
+            deletes_per_file = _read_deletes_for_task(self._io, task)
+            for batch in self._record_batches_from_scan_tasks_and_deletes([task], deletes_per_file):
+                current_batch_size = len(batch)
+                if self._limit is not None and total_row_count + current_batch_size >= self._limit:
+                    yield batch.slice(0, self._limit - total_row_count)
+                    return
+                else:
+                    yield batch
+                    total_row_count += current_batch_size
+
     def _record_batches_from_scan_tasks_and_deletes(
         self, tasks: Iterable[FileScanTask], deletes_per_file: dict[str, list[ChunkedArray]]
     ) -> Iterator[pa.RecordBatch]:

pyiceberg/table/__init__.py

@@ -2381,7 +2381,7 @@ def _to_arrow_batch_reader_via_file_scan_tasks(
         scan.case_sensitive,
         scan.limit,
         dictionary_columns=dictionary_columns,
-    ).to_record_batches(tasks)
+    ).to_record_batches_lazy(tasks)
 
     if dictionary_columns:
         # schema_to_pyarrow returns plain types, but ArrowScan yields dictionary-encoded

tests/io/test_pyarrow.py

@@ -63,7 +63,7 @@
     Or,
 )
 from pyiceberg.expressions.literals import literal
-from pyiceberg.io import S3_RETRY_STRATEGY_IMPL, S3_SSL_CA_CERT, InputStream, OutputStream, load_file_io
+from pyiceberg.io import S3_RETRY_STRATEGY_IMPL, S3_SSL_CA_CERT, FileIO, InputFile, InputStream, OutputFile, OutputStream, load_file_io
 from pyiceberg.io.pyarrow import (
     ICEBERG_SCHEMA,
     PYARROW_PARQUET_FIELD_ID_KEY,
@@ -5524,3 +5524,145 @@ def test_dictionary_columns_produces_dict_encoded_output(tmpdir: str) -> None:
 
     # Values must be identical
     assert result_plain.column("label").to_pylist() == result_dict.column("label").to_pylist()
+
+
+class _CountingFileIO(FileIO):
+    """FileIO decorator recording how many files are opened for reading."""
+
+    def __init__(self, inner: FileIO) -> None:
+        super().__init__()
+        self._inner = inner
+        self.files_opened_for_read = 0
+
+    def new_input(self, location: str) -> InputFile:
+        self.files_opened_for_read += 1
+        return self._inner.new_input(location)
+
+    def new_output(self, location: str) -> OutputFile:
+        return self._inner.new_output(location)
+
+    def delete(self, location: str | InputFile | OutputFile) -> None:
+        self._inner.delete(location)
+
+
+class _MockScan:
+    """Minimal scan stub for `_to_arrow_batch_reader_via_file_scan_tasks`."""
+
+    def __init__(self, table_metadata: TableMetadataV2, io: FileIO, limit: int | None = None) -> None:
+        self.table_metadata = table_metadata
+        self.io = io
+        self.row_filter = AlwaysTrue()
+        self.case_sensitive = True
+        self.limit = limit
+
+
+def _write_lazy_reader_test_files(
+    tmpdir: str, num_files: int, rows_per_file: int
+) -> tuple[Schema, TableMetadataV2, list[FileScanTask]]:
+    """Write `num_files` parquet files with consecutive, non-overlapping ids."""
+    iceberg_schema = Schema(
+        NestedField(1, "id", IntegerType(), required=False),
+    )
+    arrow_schema = pa.schema([pa.field("id", pa.int32(), nullable=True, metadata={PYARROW_PARQUET_FIELD_ID_KEY: "1"})])
+    table_metadata = TableMetadataV2(
+        location=f"file://{tmpdir}",
+        last_column_id=1,
+        format_version=2,
+        schemas=[iceberg_schema],
+        partition_specs=[PartitionSpec()],
+    )
+    tasks = []
+    for file_idx in range(num_files):
+        start_id = file_idx * rows_per_file
+        arrow_table = pa.table(
+            [pa.array(range(start_id, start_id + rows_per_file), type=pa.int32())],
+            schema=arrow_schema,
+        )
+        data_file = _write_table_to_data_file(f"{tmpdir}/lazy_reader_{file_idx}.parquet", arrow_schema, arrow_table)
+        data_file.spec_id = 0
+        tasks.append(FileScanTask(data_file))
+    return iceberg_schema, table_metadata, tasks
+
+
+def test_to_arrow_batch_reader_does_not_read_ahead(tmpdir: str) -> None:
+    """Regression test for https://github.com/apache/iceberg-python/issues/2407.
+
+    `to_arrow_batch_reader()` documents low-memory streaming ("a RecordBatch is
+    read one at a time"), but the reader used to fan every file scan task out to
+    the executor and materialize each file's batches into a list, so taking a
+    single batch from the reader read every file in the scan.
+    """
+    from pyiceberg.table import _to_arrow_batch_reader_via_file_scan_tasks
+
+    num_files = 4
+    iceberg_schema, table_metadata, tasks = _write_lazy_reader_test_files(tmpdir, num_files, rows_per_file=5000)
+    io = _CountingFileIO(PyArrowFileIO())
+
+    reader = _to_arrow_batch_reader_via_file_scan_tasks(_MockScan(table_metadata, io), iceberg_schema, tasks)
+
+    first_batch = next(reader)
+    assert first_batch.num_rows > 0
+    assert io.files_opened_for_read == 1, (
+        f"expected exactly 1 file opened after consuming the first batch, got {io.files_opened_for_read}"
+    )
+
+
+def test_to_record_batches_lazy_matches_eager(tmpdir: str) -> None:
+    """The lazy streaming path must return the same rows, in the same order, as the threaded path."""
+    num_files = 3
+    rows_per_file = 2500
+    total_rows = num_files * rows_per_file
+    iceberg_schema, table_metadata, tasks = _write_lazy_reader_test_files(tmpdir, num_files, rows_per_file)
+    expected_ids = list(range(total_rows))
+
+    for limit in (None, 0, 1, 100, total_rows, total_rows + 10):
+        eager_scan = ArrowScan(table_metadata, PyArrowFileIO(), iceberg_schema, AlwaysTrue(), True, limit)
+        lazy_scan = ArrowScan(table_metadata, PyArrowFileIO(), iceberg_schema, AlwaysTrue(), True, limit)
+
+        eager_ids = [row for batch in eager_scan.to_record_batches(tasks) for row in batch.column("id").to_pylist()]
+        lazy_ids = [row for batch in lazy_scan.to_record_batches_lazy(tasks) for row in batch.column("id").to_pylist()]
+
+        assert lazy_ids == eager_ids, f"limit={limit}: lazy path diverged from the threaded path"
+        assert lazy_ids == expected_ids[: len(lazy_ids)], f"limit={limit}: unexpected row contents or ordering"
+        if limit is None or limit >= total_rows:
+            assert len(lazy_ids) == total_rows
+        else:
+            assert len(lazy_ids) == limit
+
+
+def test_to_record_batches_lazy_applies_positional_deletes(tmpdir: str) -> None:
+    """The lazy path must apply per-task positional deletes exactly like the eager path."""
+    from pyiceberg.table import _to_arrow_batch_reader_via_file_scan_tasks
+
+    iceberg_schema = Schema(
+        NestedField(1, "id", IntegerType(), required=False),
+    )
+    arrow_schema = pa.schema([pa.field("id", pa.int32(), nullable=True, metadata={PYARROW_PARQUET_FIELD_ID_KEY: "1"})])
+    table_metadata = TableMetadataV2(
+        location=f"file://{tmpdir}",
+        last_column_id=1,
+        format_version=2,
+        schemas=[iceberg_schema],
+        partition_specs=[PartitionSpec()],
+    )
+
+    data_file = _write_table_to_data_file(
+        f"{tmpdir}/lazy_reader_deletes.parquet",
+        arrow_schema,
+        pa.table([pa.array([1, 2, 3, 4], type=pa.int32())], schema=arrow_schema),
+    )
+    data_file.spec_id = 0
+
+    # Positional delete of row position 2 (value 3)
+    deletes_path = f"{tmpdir}/lazy_reader_deletes_pos.parquet"
+    pq.write_table(pa.table({"file_path": [data_file.file_path], "pos": [2]}), deletes_path)
+    delete_file = DataFile.from_args(
+        content=DataFileContent.POSITION_DELETES, file_path=deletes_path, file_format=FileFormat.PARQUET
+    )
+    tasks = [FileScanTask(data_file, delete_files={delete_file})]
+
+    result = _to_arrow_batch_reader_via_file_scan_tasks(
+        _MockScan(table_metadata, PyArrowFileIO()), iceberg_schema, tasks
+    ).read_all()
+
+    assert result.column("id").to_pylist() == [1, 2, 4]