Skip to content

Commit 4ecaf15

Browse files
author
Paul Mathew
committed
feat(2152): support pa.RecordBatchReader in Table.append/overwrite
Currently `Table.append(df)` and `Table.overwrite(df)` only accept a materialised `pa.Table`, which forces callers to load the entire dataset into memory before writing. This makes pyiceberg unusable for large or unbounded inputs and has been a recurring complaint (#1004, #2152, dlt-hub#3753). Allow `pa.RecordBatchReader` as an alternative input. When a reader is provided, batches are streamed and microbatched into target-sized Parquet files via the new `bin_pack_record_batches` helper, then committed in a single snapshot via the existing fast_append path. Memory is bounded by `write.target-file-size-bytes` (default 512 MiB) per worker rather than the full input size. Scope of this PR — unpartitioned tables only. Streaming into partitioned tables raises NotImplementedError pointing back to #2152; partitioned support needs additional design (high-cardinality partition handling, per-partition rolling writers) and is tracked as a follow-up. Mirrors iceberg-go#369's staging — that project shipped unpartitioned streaming first. Internal note: the implementation buffers up to `target_file_size` of in- memory RecordBatches before flushing to a Parquet file. A more memory- efficient rolling-ParquetWriter approach is a planned follow-up that will benefit from the `OutputStream.tell()` API added in #2998.
1 parent d008a04 commit 4ecaf15

6 files changed

Lines changed: 460 additions & 29 deletions

File tree

‎mkdocs/docs/api.md‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -365,6 +365,17 @@ for buf in tbl.scan().to_arrow_batch_reader():
365365
print(f"Buffer contains {len(buf)} rows")
366366
```
367367

368+
### Streaming writes from a `RecordBatchReader`
369+
370+
`tbl.append()` and `tbl.overwrite()` also accept a `pyarrow.RecordBatchReader` directly, which lets you write datasets that don't fit in memory without materialising them as a `pa.Table` first. PyIceberg consumes the reader once and microbatches it into Parquet files of approximately `write.target-file-size-bytes` (default 512 MiB), keeping memory usage bounded by the target size. All files are committed in a single snapshot.
371+
372+
```python
373+
reader = pa.RecordBatchReader.from_batches(schema, batch_iter)
374+
tbl.append(reader)
375+
```
376+
377+
Streaming writes are currently only supported on **unpartitioned** tables. For a partitioned table, materialise the reader as a `pa.Table` first, or follow [#2152](https://github.com/apache/iceberg-python/issues/2152) for the partitioned support tracked as a follow-up.
378+
368379
To avoid any type inconsistencies during writing, you can convert the Iceberg table schema to Arrow:
369380

370381
```python

‎pyiceberg/io/pyarrow.py‎

Lines changed: 55 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2681,6 +2681,32 @@ def bin_pack_arrow_table(tbl: pa.Table, target_file_size: int) -> Iterator[list[
26812681
return bin_packed_record_batches
26822682

26832683

2684+
def bin_pack_record_batches(batches: Iterable[pa.RecordBatch], target_file_size: int) -> Iterator[list[pa.RecordBatch]]:
2685+
"""Microbatch a single-pass stream of RecordBatches into target-sized groups.
2686+
2687+
Unlike :func:`bin_pack_arrow_table`, this consumes ``batches`` lazily and
2688+
holds at most one in-flight buffer in memory, bounded by ``target_file_size``.
2689+
Suitable for streaming inputs (``pa.RecordBatchReader``,
2690+
``Iterator[pa.RecordBatch]``) where the total size is unknown up front and
2691+
the caller cannot afford to materialise the full dataset.
2692+
2693+
Each yielded list of batches is intended to be written as a single Parquet
2694+
data file. Because this is single-pass FIFO accumulation (no lookback), the
2695+
last bin may be smaller than ``target_file_size``.
2696+
"""
2697+
buffer: list[pa.RecordBatch] = []
2698+
buffer_bytes = 0
2699+
for batch in batches:
2700+
buffer.append(batch)
2701+
buffer_bytes += batch.nbytes
2702+
if buffer_bytes >= target_file_size:
2703+
yield buffer
2704+
buffer = []
2705+
buffer_bytes = 0
2706+
if buffer:
2707+
yield buffer
2708+
2709+
26842710
def _check_pyarrow_schema_compatible(
26852711
requested_schema: Schema,
26862712
provided_schema: pa.Schema,
@@ -2800,15 +2826,24 @@ def _get_parquet_writer_kwargs(table_properties: Properties) -> dict[str, Any]:
28002826

28012827
def _dataframe_to_data_files(
28022828
table_metadata: TableMetadata,
2803-
df: pa.Table,
2829+
df: pa.Table | pa.RecordBatchReader,
28042830
io: FileIO,
28052831
write_uuid: uuid.UUID | None = None,
28062832
counter: itertools.count[int] | None = None,
28072833
) -> Iterable[DataFile]:
2808-
"""Convert a PyArrow table into a DataFile.
2834+
"""Convert a PyArrow Table or RecordBatchReader into DataFiles.
2835+
2836+
For a ``pa.Table`` the data is materialised in memory and bin-packed into
2837+
target-sized files (with partition splitting if the table is partitioned).
2838+
2839+
For a ``pa.RecordBatchReader`` batches are streamed and microbatched into
2840+
target-sized files using bounded memory (see :func:`bin_pack_record_batches`).
2841+
Streaming writes are currently only supported on unpartitioned tables;
2842+
partitioned support is tracked in
2843+
https://github.com/apache/iceberg-python/issues/2152.
28092844
28102845
Returns:
2811-
An iterable that supplies datafiles that represent the table.
2846+
An iterable that supplies datafiles that represent the input data.
28122847
"""
28132848
from pyiceberg.table import DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE, TableProperties, WriteTask
28142849

@@ -2828,6 +2863,23 @@ def _dataframe_to_data_files(
28282863
format_version=table_metadata.format_version,
28292864
)
28302865

2866+
if isinstance(df, pa.RecordBatchReader):
2867+
if not table_metadata.spec().is_unpartitioned():
2868+
raise NotImplementedError(
2869+
"Writing a pa.RecordBatchReader to a partitioned table is not yet supported. "
2870+
"Materialise the reader as a pa.Table first, or follow "
2871+
"https://github.com/apache/iceberg-python/issues/2152 for partitioned streaming support."
2872+
)
2873+
yield from write_file(
2874+
io=io,
2875+
table_metadata=table_metadata,
2876+
tasks=(
2877+
WriteTask(write_uuid=write_uuid, task_id=next(counter), record_batches=batches, schema=task_schema)
2878+
for batches in bin_pack_record_batches(df, target_file_size)
2879+
),
2880+
)
2881+
return
2882+
28312883
if table_metadata.spec().is_unpartitioned():
28322884
yield from write_file(
28332885
io=io,

‎pyiceberg/table/__init__.py‎

Lines changed: 100 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -450,12 +450,45 @@ def update_statistics(self) -> UpdateStatistics:
450450
"""
451451
return UpdateStatistics(transaction=self)
452452

453-
def append(self, df: pa.Table, snapshot_properties: dict[str, str] = EMPTY_DICT, branch: str | None = MAIN_BRANCH) -> None:
453+
def append(
454+
self,
455+
df: pa.Table | pa.RecordBatchReader,
456+
snapshot_properties: dict[str, str] = EMPTY_DICT,
457+
branch: str | None = MAIN_BRANCH,
458+
) -> None:
454459
"""
455-
Shorthand API for appending a PyArrow table to a table transaction.
460+
Shorthand API for appending PyArrow data to a table transaction.
461+
462+
Accepts either a fully materialised ``pa.Table`` or a streaming
463+
``pa.RecordBatchReader``. Streaming is microbatched by
464+
``write.target-file-size-bytes`` so memory stays bounded; the reader is
465+
consumed once and cannot be reused.
466+
467+
Streaming writes are currently only supported on unpartitioned tables;
468+
passing a ``pa.RecordBatchReader`` for a partitioned table raises
469+
``NotImplementedError``. See
470+
https://github.com/apache/iceberg-python/issues/2152.
471+
472+
Note:
473+
When ``df`` is a ``pa.RecordBatchReader`` the reader is consumed
474+
once and cannot be replayed. If the catalog commit fails (e.g.
475+
``CommitFailedException`` from a concurrent writer) the reader is
476+
already drained and a naive retry will append zero rows. Callers
477+
that need at-least-once semantics should either:
478+
479+
- reconstruct the reader on each attempt via a factory callable,
480+
or
481+
- use a two-stage pattern — write Parquet files explicitly and
482+
then call :meth:`add_files` (whose input is a replayable list of
483+
paths) within a retry loop.
484+
485+
Failures during the write stage (mid-stream reader exception, S3
486+
errors) do not commit a snapshot, but may leave orphan data files
487+
in storage that are not referenced by any snapshot. Clean these
488+
up with expire/orphan-file maintenance jobs.
456489
457490
Args:
458-
df: The Arrow dataframe that will be appended to overwrite the table
491+
df: An Arrow Table or a RecordBatchReader of records to append.
459492
snapshot_properties: Custom properties to be added to the snapshot summary
460493
branch: Branch Reference to run the append operation
461494
"""
@@ -466,8 +499,8 @@ def append(self, df: pa.Table, snapshot_properties: dict[str, str] = EMPTY_DICT,
466499

467500
from pyiceberg.io.pyarrow import _check_pyarrow_schema_compatible, _dataframe_to_data_files
468501

469-
if not isinstance(df, pa.Table):
470-
raise ValueError(f"Expected PyArrow table, got: {df}")
502+
if not isinstance(df, (pa.Table, pa.RecordBatchReader)):
503+
raise ValueError(f"Expected pa.Table or pa.RecordBatchReader, got: {df}")
471504

472505
downcast_ns_timestamp_to_us = Config().get_bool(DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE) or False
473506
_check_pyarrow_schema_compatible(
@@ -478,12 +511,14 @@ def append(self, df: pa.Table, snapshot_properties: dict[str, str] = EMPTY_DICT,
478511
)
479512

480513
with self._append_snapshot_producer(snapshot_properties, branch=branch) as append_files:
481-
# skip writing data files if the dataframe is empty
482-
if df.shape[0] > 0:
483-
data_files = list(
484-
_dataframe_to_data_files(
485-
table_metadata=self.table_metadata, write_uuid=append_files.commit_uuid, df=df, io=self._table.io
486-
)
514+
# For pa.Table we can short-circuit empty inputs cheaply. For a
515+
# RecordBatchReader the stream is consumed lazily by
516+
# _dataframe_to_data_files and an empty reader simply yields zero
517+
# data files (the snapshot is still committed for symmetry with the
518+
# pa.Table case where empty inputs also produce a snapshot).
519+
if isinstance(df, pa.RecordBatchReader) or df.shape[0] > 0:
520+
data_files = _dataframe_to_data_files(
521+
table_metadata=self.table_metadata, write_uuid=append_files.commit_uuid, df=df, io=self._table.io
487522
)
488523
for data_file in data_files:
489524
append_files.append_data_file(data_file)
@@ -555,14 +590,42 @@ def dynamic_partition_overwrite(
555590

556591
def overwrite(
557592
self,
558-
df: pa.Table,
593+
df: pa.Table | pa.RecordBatchReader,
559594
overwrite_filter: BooleanExpression | str = ALWAYS_TRUE,
560595
snapshot_properties: dict[str, str] = EMPTY_DICT,
561596
case_sensitive: bool = True,
562597
branch: str | None = MAIN_BRANCH,
563598
) -> None:
564599
"""
565-
Shorthand for adding a table overwrite with a PyArrow table to the transaction.
600+
Shorthand for adding a table overwrite with a PyArrow table or RecordBatchReader to the transaction.
601+
602+
Accepts either a fully materialised ``pa.Table`` or a streaming
603+
``pa.RecordBatchReader``. Streaming is microbatched by
604+
``write.target-file-size-bytes`` so memory stays bounded; the reader is
605+
consumed once and cannot be reused.
606+
607+
Streaming writes are currently only supported on unpartitioned tables;
608+
passing a ``pa.RecordBatchReader`` for a partitioned table raises
609+
``NotImplementedError``. See
610+
https://github.com/apache/iceberg-python/issues/2152.
611+
612+
Note:
613+
When ``df`` is a ``pa.RecordBatchReader`` the reader is consumed
614+
once and cannot be replayed. If the catalog commit fails (e.g.
615+
``CommitFailedException`` from a concurrent writer) the reader is
616+
already drained and a naive retry will write zero rows. Callers
617+
that need at-least-once semantics should either:
618+
619+
- reconstruct the reader on each attempt via a factory callable,
620+
or
621+
- use a two-stage pattern — write Parquet files explicitly and
622+
then call :meth:`add_files` (whose input is a replayable list
623+
of paths) within a retry loop.
624+
625+
Failures during the write stage (mid-stream reader exception, S3
626+
errors) do not commit a snapshot, but may leave orphan data files
627+
in storage that are not referenced by any snapshot. Clean these
628+
up with expire/orphan-file maintenance jobs.
566629
567630
An overwrite may produce zero or more snapshots based on the operation:
568631
@@ -571,7 +634,7 @@ def overwrite(
571634
- APPEND: In case new data is being inserted into the table.
572635
573636
Args:
574-
df: The Arrow dataframe that will be used to overwrite the table
637+
df: An Arrow Table or a RecordBatchReader of records to write.
575638
overwrite_filter: ALWAYS_TRUE when you overwrite all the data,
576639
or a boolean expression in case of a partial overwrite
577640
snapshot_properties: Custom properties to be added to the snapshot summary
@@ -585,8 +648,8 @@ def overwrite(
585648

586649
from pyiceberg.io.pyarrow import _check_pyarrow_schema_compatible, _dataframe_to_data_files
587650

588-
if not isinstance(df, pa.Table):
589-
raise ValueError(f"Expected PyArrow table, got: {df}")
651+
if not isinstance(df, (pa.Table, pa.RecordBatchReader)):
652+
raise ValueError(f"Expected pa.Table or pa.RecordBatchReader, got: {df}")
590653

591654
downcast_ns_timestamp_to_us = Config().get_bool(DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE) or False
592655
_check_pyarrow_schema_compatible(
@@ -606,8 +669,8 @@ def overwrite(
606669
)
607670

608671
with self._append_snapshot_producer(snapshot_properties, branch=branch) as append_files:
609-
# skip writing data files if the dataframe is empty
610-
if df.shape[0] > 0:
672+
# See append() for the empty-input handling rationale.
673+
if isinstance(df, pa.RecordBatchReader) or df.shape[0] > 0:
611674
data_files = _dataframe_to_data_files(
612675
table_metadata=self.table_metadata, write_uuid=append_files.commit_uuid, df=df, io=self._table.io
613676
)
@@ -1373,12 +1436,21 @@ def upsert(
13731436
snapshot_properties=snapshot_properties,
13741437
)
13751438

1376-
def append(self, df: pa.Table, snapshot_properties: dict[str, str] = EMPTY_DICT, branch: str | None = MAIN_BRANCH) -> None:
1439+
def append(
1440+
self,
1441+
df: pa.Table | pa.RecordBatchReader,
1442+
snapshot_properties: dict[str, str] = EMPTY_DICT,
1443+
branch: str | None = MAIN_BRANCH,
1444+
) -> None:
13771445
"""
1378-
Shorthand API for appending a PyArrow table to the table.
1446+
Shorthand API for appending PyArrow data to the table.
1447+
1448+
Accepts either a ``pa.Table`` or a streaming ``pa.RecordBatchReader``.
1449+
See :meth:`Transaction.append` for streaming semantics and partition
1450+
limitations.
13791451
13801452
Args:
1381-
df: The Arrow dataframe that will be appended to overwrite the table
1453+
df: An Arrow Table or a RecordBatchReader of records to append.
13821454
snapshot_properties: Custom properties to be added to the snapshot summary
13831455
branch: Branch Reference to run the append operation
13841456
"""
@@ -1401,14 +1473,18 @@ def dynamic_partition_overwrite(
14011473

14021474
def overwrite(
14031475
self,
1404-
df: pa.Table,
1476+
df: pa.Table | pa.RecordBatchReader,
14051477
overwrite_filter: BooleanExpression | str = ALWAYS_TRUE,
14061478
snapshot_properties: dict[str, str] = EMPTY_DICT,
14071479
case_sensitive: bool = True,
14081480
branch: str | None = MAIN_BRANCH,
14091481
) -> None:
14101482
"""
1411-
Shorthand for overwriting the table with a PyArrow table.
1483+
Shorthand for overwriting the table with a PyArrow Table or RecordBatchReader.
1484+
1485+
Accepts either a ``pa.Table`` or a streaming ``pa.RecordBatchReader``.
1486+
See :meth:`Transaction.overwrite` for streaming semantics and partition
1487+
limitations.
14121488
14131489
An overwrite may produce zero or more snapshots based on the operation:
14141490
@@ -1417,7 +1493,7 @@ def overwrite(
14171493
- APPEND: In case new data is being inserted into the table.
14181494
14191495
Args:
1420-
df: The Arrow dataframe that will be used to overwrite the table
1496+
df: An Arrow Table or a RecordBatchReader of records to write.
14211497
overwrite_filter: ALWAYS_TRUE when you overwrite all the data,
14221498
or a boolean expression in case of a partial overwrite
14231499
snapshot_properties: Custom properties to be added to the snapshot summary

0 commit comments

Comments
 (0)