Skip to content

fix(io): stream RecordBatchReader writes through a rolling ParquetWriter #3661

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

Open
Sushankthatipally wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(io): stream RecordBatchReader writes through a rolling ParquetWriter #3661

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 1 addition & 1 deletion mkdocs/docs/api.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -367,7 +367,7 @@ for buf in tbl.scan().to_arrow_batch_reader():

### Streaming writes from a `RecordBatchReader`

`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.
`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, writing batches through a rolling Parquet writer that rolls a new file each time the on-disk file size hits `write.target-file-size-bytes` (default 512 MiB). Each input `RecordBatch` becomes a Parquet row group, capped at `write.parquet.row-group-limit` rows (default 1M) — caller batch size sets the lower bound on row group size, the property enforces the upper bound. All files are committed in a single snapshot.

```python
reader = pa.RecordBatchReader.from_batches(schema, batch_iter)
Expand Down
182 changes: 138 additions & 44 deletions pyiceberg/io/pyarrow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2715,39 +2715,136 @@ def bin_pack_arrow_table(tbl: pa.Table, target_file_size: int) -> Iterator[list[
return bin_packed_record_batches


def bin_pack_record_batches(batches: Iterable[pa.RecordBatch], target_file_size: int) -> Iterator[list[pa.RecordBatch]]:
"""Microbatch a single-pass stream of RecordBatches into target-sized groups.
def _record_batches_to_data_files(
table_metadata: TableMetadata,
reader: pa.RecordBatchReader,
io: FileIO,
write_uuid: uuid.UUID | None = None,
counter: itertools.count[int] | None = None,
) -> Iterator[DataFile]:
"""Stream a ``pa.RecordBatchReader`` into Parquet data files via a rolling ``pq.ParquetWriter``.

Each input ``RecordBatch`` is written directly via ``writer.write_batch``. File
rollover is driven by ``OutputStream.tell()``: after each batch, if
``tell() >= write.target-file-size-bytes`` the current writer is closed
(footer written) and a new file is opened. The threshold is measured in
compressed on-disk bytes, matching the spec-defined semantics of
``write.target-file-size-bytes`` and the Java/Spark/Flink writers.

Row groups are capped at ``write.parquet.row-group-limit`` rows (default 1M)
via the ``row_group_size`` argument to ``write_batch``: a batch larger than
the cap is split into multiple row groups, each <= the cap; a smaller batch
becomes a single row group of its own size. Callers control the lower bound
of row group size via their choice of input batch size; this enforces the
upper bound, matching the materialised ``pa.Table`` write path.

Peak memory is bounded by one input ``RecordBatch`` plus the
``ParquetWriter``'s internal page buffer, regardless of dataset size,
``target_file_size``, or number of files produced.

Streaming writes to partitioned tables are not yet supported - see
https://github.com/apache/iceberg-python/issues/2152.
"""
from pyiceberg.table import DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE, TableProperties

Unlike :func:`bin_pack_arrow_table`, this consumes ``batches`` lazily and
holds at most one in-flight buffer in memory, bounded by ``target_file_size``.
Suitable for streaming inputs (``pa.RecordBatchReader``,
``Iterator[pa.RecordBatch]``) where the total size is unknown up front and
the caller cannot afford to materialise the full dataset.
if not table_metadata.spec().is_unpartitioned():
raise NotImplementedError(
"Writing a pa.RecordBatchReader to a partitioned table is not yet supported. "
"Materialise the reader as a pa.Table first, or follow "
"https://github.com/apache/iceberg-python/issues/2152 for partitioned streaming support."
)

Each yielded list of batches is intended to be written as a single Parquet
data file. Because this is single-pass FIFO accumulation (no lookback), the
last bin may be smaller than ``target_file_size``.
counter = counter or itertools.count(0)
write_uuid = write_uuid or uuid.uuid4()
target_file_size: int = property_as_int( # type: ignore # The property is set with non-None value.
properties=table_metadata.properties,
property_name=TableProperties.WRITE_TARGET_FILE_SIZE_BYTES,
default=TableProperties.WRITE_TARGET_FILE_SIZE_BYTES_DEFAULT,
)
name_mapping = table_metadata.schema().name_mapping
downcast_ns_timestamp_to_us = Config().get_bool(DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE) or False
task_schema = pyarrow_to_schema(
reader.schema,
name_mapping=name_mapping,
downcast_ns_timestamp_to_us=downcast_ns_timestamp_to_us,
format_version=table_metadata.format_version,
)
table_schema = table_metadata.schema()
if (sanitized_schema := sanitize_column_names(table_schema)) != table_schema:
file_schema = sanitized_schema
else:
file_schema = table_schema

Note:
``target_file_size`` is measured in **uncompressed in-memory** Arrow
bytes (``RecordBatch.nbytes``), not compressed on-disk Parquet bytes.
The resulting Parquet file after compression is typically 3-10×
smaller than ``target_file_size``. Matches the existing
:func:`bin_pack_arrow_table` semantics; both will be tightened to true
on-disk bytes once the writer is switched to a rolling-
``ParquetWriter`` with ``OutputStream.tell()`` (#2998).
"""
buffer: list[pa.RecordBatch] = []
buffer_bytes = 0
for batch in batches:
buffer.append(batch)
buffer_bytes += batch.nbytes
if buffer_bytes >= target_file_size:
yield buffer
buffer = []
buffer_bytes = 0
if buffer:
yield buffer
parquet_writer_kwargs = _get_parquet_writer_kwargs(table_metadata.properties)
row_group_size = property_as_int(
properties=table_metadata.properties,
property_name=TableProperties.PARQUET_ROW_GROUP_LIMIT,
default=TableProperties.PARQUET_ROW_GROUP_LIMIT_DEFAULT,
)
location_provider = load_location_provider(table_location=table_metadata.location, table_properties=table_metadata.properties)
stats_columns = compute_statistics_plan(file_schema, table_metadata.properties)
column_mapping = parquet_path_to_id_mapping(file_schema)

def _transform(batch: pa.RecordBatch) -> pa.RecordBatch:
return _to_requested_schema(
requested_schema=file_schema,
file_schema=task_schema,
batch=batch,
downcast_ns_timestamp_to_us=downcast_ns_timestamp_to_us,
include_field_ids=True,
)

def _new_data_file_path() -> str:
# Mirrors WriteTask.generate_data_file_filename to keep file names
# compatible with the materialised write path.
filename = f"00000-{next(counter)}-{write_uuid}.parquet"
return location_provider.new_data_location(data_file_name=filename)

batches = iter(reader)
while True:
# Pull the next batch up front. If the reader is exhausted (either at the
# very start or between rolled files), we're done - yield nothing further.
try:
first_batch = next(batches)
except StopIteration:
return

transformed_first = _transform(first_batch)
file_path = _new_data_file_path()
fo = io.new_output(file_path)
with fo.create(overwrite=True) as fos:
with pq.ParquetWriter(
fos, schema=transformed_first.schema, store_decimal_as_integer=True, **parquet_writer_kwargs
) as writer:
writer.write_batch(transformed_first, row_group_size=row_group_size)
# Keep writing into this file until the on-disk byte threshold is
# crossed. tell() advances as write_batch flushes encoded pages to
# the stream - files end up close to but slightly above
# target_file_size (lag bounded by one Parquet data page).
while fos.tell() < target_file_size:
try:
batch = next(batches)
except StopIteration:
break
writer.write_batch(_transform(batch), row_group_size=row_group_size)

statistics = data_file_statistics_from_parquet_metadata(
parquet_metadata=writer.writer.metadata,
stats_columns=stats_columns,
parquet_column_mapping=column_mapping,
)
yield DataFile.from_args(
content=DataFileContent.DATA,
file_path=file_path,
file_format=FileFormat.PARQUET,
partition=Record(),
file_size_in_bytes=len(fo),
sort_order_id=None,
spec_id=table_metadata.default_spec_id,
equality_ids=None,
key_metadata=None,
**statistics.to_serialized_dict(),
)


def _check_pyarrow_schema_compatible(
Expand Down Expand Up @@ -2879,8 +2976,8 @@ def _dataframe_to_data_files(
For a ``pa.Table`` the data is materialised in memory and bin-packed into
target-sized files (with partition splitting if the table is partitioned).

For a ``pa.RecordBatchReader`` batches are streamed and microbatched into
target-sized files using bounded memory (see :func:`bin_pack_record_batches`).
For a ``pa.RecordBatchReader`` batches are streamed through a rolling
``pq.ParquetWriter`` with constant memory (see :func:`_record_batches_to_data_files`).
Streaming writes are currently only supported on unpartitioned tables;
partitioned support is tracked in
https://github.com/apache/iceberg-python/issues/2152.
Expand All @@ -2907,19 +3004,16 @@ def _dataframe_to_data_files(
)

if isinstance(df, pa.RecordBatchReader):
if not table_metadata.spec().is_unpartitioned():
raise NotImplementedError(
"Writing a pa.RecordBatchReader to a partitioned table is not yet supported. "
"Materialise the reader as a pa.Table first, or follow "
"https://github.com/apache/iceberg-python/issues/2152 for partitioned streaming support."
)
yield from write_file(
io=io,
# Streaming path: rolling ParquetWriter driven by OutputStream.tell() for
# constant-memory writes and on-disk-accurate file sizes. The
# partitioned-table NotImplementedError is raised inside
# _record_batches_to_data_files.
yield from _record_batches_to_data_files(
table_metadata=table_metadata,
tasks=(
WriteTask(write_uuid=write_uuid, task_id=next(counter), record_batches=batches, schema=task_schema)
for batches in bin_pack_record_batches(df, target_file_size)
),
reader=df,
io=io,
write_uuid=write_uuid,
counter=counter,
)
return

Expand Down
41 changes: 21 additions & 20 deletions pyiceberg/table/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -467,8 +467,9 @@ def append(
Shorthand API for appending PyArrow data to a table transaction.

Accepts either a fully materialised ``pa.Table`` or a streaming
``pa.RecordBatchReader``. Streaming is microbatched by
``write.target-file-size-bytes`` so memory stays bounded; the reader is
``pa.RecordBatchReader``. For a reader, batches are written through a
rolling ``pq.ParquetWriter`` and a new file is rolled each time the
on-disk file size hits ``write.target-file-size-bytes``. The reader is
consumed once and cannot be reused.

Streaming writes are currently only supported on unpartitioned tables;
Expand All @@ -494,13 +495,12 @@ def append(
in storage that are not referenced by any snapshot. Clean these
up with expire/orphan-file maintenance jobs.

``write.target-file-size-bytes`` is currently interpreted as
uncompressed in-memory Arrow bytes (the bin-packing weight) rather
than compressed on-disk Parquet bytes. The resulting files are
typically 3-10× smaller than the property suggests after
compression. This matches the existing ``pa.Table`` write path and
will be tightened once the writer is switched to a
rolling-``ParquetWriter`` with ``OutputStream.tell()`` (#2998).
For streaming inputs (``pa.RecordBatchReader``) each input
``RecordBatch`` becomes one Parquet row group. The
``write.parquet.row-group-limit`` property (rows, default 1M)
caps row group size — batches larger than the cap are split,
smaller batches are not combined. Caller batch size sets the
lower bound; pyiceberg enforces the upper bound.

Args:
df: An Arrow Table or a RecordBatchReader of records to append.
Expand Down Expand Up @@ -615,8 +615,9 @@ def overwrite(
Shorthand for adding a table overwrite with a PyArrow table or RecordBatchReader to the transaction.

Accepts either a fully materialised ``pa.Table`` or a streaming
``pa.RecordBatchReader``. Streaming is microbatched by
``write.target-file-size-bytes`` so memory stays bounded; the reader is
``pa.RecordBatchReader``. For a reader, batches are written through a
rolling ``pq.ParquetWriter`` and a new file is rolled each time the
on-disk file size hits ``write.target-file-size-bytes``. The reader is
consumed once and cannot be reused.

Streaming writes are currently only supported on unpartitioned tables;
Expand All @@ -642,13 +643,12 @@ def overwrite(
in storage that are not referenced by any snapshot. Clean these
up with expire/orphan-file maintenance jobs.

``write.target-file-size-bytes`` is currently interpreted as
uncompressed in-memory Arrow bytes (the bin-packing weight) rather
than compressed on-disk Parquet bytes. The resulting files are
typically 3-10× smaller than the property suggests after
compression. This matches the existing ``pa.Table`` write path and
will be tightened once the writer is switched to a
rolling-``ParquetWriter`` with ``OutputStream.tell()`` (#2998).
For streaming inputs (``pa.RecordBatchReader``) each input
``RecordBatch`` becomes one Parquet row group. The
``write.parquet.row-group-limit`` property (rows, default 1M)
caps row group size — batches larger than the cap are split,
smaller batches are not combined. Caller batch size sets the
lower bound; pyiceberg enforces the upper bound.

An overwrite may produce zero or more snapshots based on the operation:

Expand Down Expand Up @@ -2471,8 +2471,9 @@ def plan_files(self) -> Iterable[FileScanTask]:
options=self.options,
).plan_files(
manifests=manifests,
manifest_entry_filter=lambda manifest_entry: manifest_entry.snapshot_id in append_snapshot_ids
and manifest_entry.status == ManifestEntryStatus.ADDED,
manifest_entry_filter=lambda manifest_entry: (
manifest_entry.snapshot_id in append_snapshot_ids and manifest_entry.status == ManifestEntryStatus.ADDED
),
)

def to_arrow(self) -> pa.Table:
Expand Down
Loading