diff --git a/mkdocs/docs/api.md b/mkdocs/docs/api.md index b23f7b976e..23eeb84160 100644 --- a/mkdocs/docs/api.md +++ b/mkdocs/docs/api.md @@ -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) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index 7e1babd58d..f4f5e23f74 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -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( @@ -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. @@ -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 diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 63b87d290e..6fe162d47e 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -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; @@ -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. @@ -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; @@ -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: @@ -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: diff --git a/tests/catalog/test_catalog_behaviors.py b/tests/catalog/test_catalog_behaviors.py index b859e2d541..f00ad69dc8 100644 --- a/tests/catalog/test_catalog_behaviors.py +++ b/tests/catalog/test_catalog_behaviors.py @@ -1215,8 +1215,10 @@ def test_drop_namespace_raises_error_when_namespace_not_empty( # RecordBatchReader streaming append/overwrite tests # -# Streaming writes accept a pa.RecordBatchReader and microbatch it into target-sized -# Parquet files instead of materialising the full Arrow Table in memory. Tracks +# Streaming writes accept a pa.RecordBatchReader and write it through a rolling +# Parquet writer (row groups flushed at write.parquet.row-group-limit, files +# rolled at write.target-file-size-bytes via OutputStream.tell()) instead of +# materialising the full Arrow Table in memory. Tracks # https://github.com/apache/iceberg-python/issues/2152. @@ -1248,11 +1250,15 @@ def test_append_record_batch_reader(catalog: Catalog) -> None: def test_append_record_batch_reader_microbatched(catalog: Catalog) -> None: """A reader bigger than the per-file target produces multiple Parquet files - in a single snapshot — verifying the byte-budget microbatching path.""" + in a single snapshot — verifies file rollover via ``OutputStream.tell()``. + + Sets a tiny ``target-file-size-bytes`` so each batch's flush rolls a new + file. Each input ``RecordBatch`` is its own row group, so ``tell()`` + advances after every ``write_batch``. + """ catalog.create_namespace("default") identifier = f"default.append_record_batch_reader_microbatch_{catalog.name}" reader, total_rows = _simple_record_batch_reader(num_batches=8) - # Force every batch to roll a new file by setting an absurdly small target size. tbl = catalog.create_table( identifier=identifier, schema=reader.schema, @@ -1269,6 +1275,111 @@ def test_append_record_batch_reader_microbatched(catalog: Catalog) -> None: assert len(tbl.scan().to_arrow()) == total_rows +def test_append_record_batch_reader_row_group_limit_is_cap(catalog: Catalog) -> None: + """``write.parquet.row-group-limit`` caps the maximum rows per Parquet + row group. A single input batch larger than the cap is split into + multiple row groups, each <= the cap. The streaming path enforces the + upper bound; callers control the lower bound by their choice of input + batch size. + """ + import pyarrow.parquet as pq + + catalog.create_namespace("default") + identifier = f"default.append_record_batch_reader_row_group_limit_cap_{catalog.name}" + + row_group_cap = 250 + total_rows = 1000 # 4x the cap + schema = pa.schema([("id", pa.int64())]) + # One big batch — pyiceberg should split it into ceil(1000 / 250) = 4 row + # groups of exactly 250 rows each. + big_batch = pa.RecordBatch.from_pylist( + [{"id": i} for i in range(total_rows)], + schema=schema, + ) + reader = pa.RecordBatchReader.from_batches(schema, iter([big_batch])) + + tbl = catalog.create_table( + identifier=identifier, + schema=schema, + properties={ + TableProperties.PARQUET_ROW_GROUP_LIMIT: str(row_group_cap), + # Big enough that everything fits in one file; we're testing row + # group size, not file rollover. + TableProperties.WRITE_TARGET_FILE_SIZE_BYTES: str(64 * 1024 * 1024), + }, + ) + tbl.append(reader) + + assert len(tbl.scan().to_arrow()) == total_rows + + files = tbl.inspect.files().select(["file_path"]).to_pylist() + assert len(files) == 1, files + + file_path = files[0]["file_path"] + metadata = pq.read_metadata(tbl.io.new_input(file_path).open()) + row_group_sizes = [metadata.row_group(i).num_rows for i in range(metadata.num_row_groups)] + + # Expect 4 row groups of exactly row_group_cap rows each. Without the cap + # passed to write_batch, the whole 1000-row batch would become one row + # group — the test would fail loudly. + assert metadata.num_row_groups == total_rows // row_group_cap, row_group_sizes + for rg_size in row_group_sizes: + assert rg_size == row_group_cap, row_group_sizes + + +def test_append_record_batch_reader_target_file_size_is_on_disk_bytes(catalog: Catalog) -> None: + """The streaming write path interprets ``write.target-file-size-bytes`` as + actual on-disk compressed Parquet bytes (via ``OutputStream.tell()``), not + uncompressed in-memory Arrow bytes. This test sets a small file target, + streams several batches, and asserts each rolled file is close to the + target size — proving the spec-correct semantics. + """ + catalog.create_namespace("default") + identifier = f"default.append_record_batch_reader_target_size_{catalog.name}" + + target_bytes = 32 * 1024 # 32 KiB target — small so we get multiple files quickly + schema = pa.schema([("id", pa.int64()), ("payload", pa.large_string())]) + # ~80 bytes per row uncompressed; with zstd ~10x compression we expect + # approximately 4000 rows per ~32 KiB file. + rows_per_batch = 1000 + total_batches = 12 + batches = [ + pa.RecordBatch.from_pylist( + [{"id": i * rows_per_batch + j, "payload": f"row_{i * rows_per_batch + j:08d}"} for j in range(rows_per_batch)], + schema=schema, + ) + for i in range(total_batches) + ] + reader = pa.RecordBatchReader.from_batches(schema, iter(batches)) + expected_rows = total_batches * rows_per_batch + + tbl = catalog.create_table( + identifier=identifier, + schema=schema, + properties={TableProperties.WRITE_TARGET_FILE_SIZE_BYTES: str(target_bytes)}, + ) + tbl.append(reader) + + assert len(tbl.scan().to_arrow()) == expected_rows + + snapshot = tbl.metadata.current_snapshot() + assert snapshot is not None and snapshot.summary is not None + added_files = int(snapshot.summary["added-data-files"]) # type: ignore[arg-type] + assert added_files >= 2, snapshot.summary + + # Per-file size: every rolled file (i.e. all but possibly the last) should be + # *close to* target_bytes. The lag between tell() and write_batch is bounded + # by one Parquet data page (~1 MiB by default), so files end up slightly + # above target. We assert each rolled file is between 0.5x and 5x the + # target — a loose bound that catches the old uncompressed-Arrow-bytes + # behaviour (where files would be ~3-10x SMALLER than target). + files = tbl.inspect.files().select(["file_path", "file_size_in_bytes"]).to_pylist() + rolled_files = files[:-1] if len(files) > 1 else files + for f in rolled_files: + size = f["file_size_in_bytes"] + assert target_bytes // 2 <= size <= target_bytes * 5, f"{f['file_path']}: {size} bytes (target {target_bytes})" + + def test_append_record_batch_reader_empty(catalog: Catalog) -> None: catalog.create_namespace("default") identifier = f"default.append_record_batch_reader_empty_{catalog.name}" diff --git a/tests/integration/test_writes/test_writes.py b/tests/integration/test_writes/test_writes.py index 30fdd76ab7..974436c6e5 100644 --- a/tests/integration/test_writes/test_writes.py +++ b/tests/integration/test_writes/test_writes.py @@ -2628,7 +2628,7 @@ def test_append_record_batch_reader_multifile( spark: SparkSession, session_catalog: Catalog, arrow_table_with_null: pa.Table, format_version: int ) -> None: """Forcing a tiny target file size should produce >1 data file in a single - snapshot, proving the byte-budget rollover in bin_pack_record_batches fires + snapshot, proving the rolling ParquetWriter's tell()-based rollover fires end-to-end and the resulting files are valid Iceberg data files (Spark reads them all).""" identifier = f"default.streaming_append_multifile_v{format_version}" diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index 532311899d..0d72134b2f 100644 --- a/tests/io/test_pyarrow.py +++ b/tests/io/test_pyarrow.py @@ -20,7 +20,6 @@ import tempfile import uuid import warnings -from collections.abc import Iterator from datetime import date, datetime, timezone from pathlib import Path from typing import Any @@ -77,7 +76,6 @@ _task_to_record_batches, _to_requested_schema, bin_pack_arrow_table, - bin_pack_record_batches, compute_statistics_plan, data_file_statistics_from_parquet_metadata, expression_to_pyarrow, @@ -2435,50 +2433,6 @@ def test_bin_pack_arrow_table_target_size_smaller_than_row(arrow_table_with_null assert sum(batch.num_rows for bin_ in bin_packed for batch in bin_) == arrow_table_with_null.num_rows -def test_bin_pack_record_batches_single_bin(arrow_table_with_null: pa.Table) -> None: - batches = arrow_table_with_null.to_batches() - bins = list(bin_pack_record_batches(iter(batches), target_file_size=arrow_table_with_null.nbytes * 10)) - # everything fits in one bin - assert len(bins) == 1 - assert sum(b.num_rows for b in bins[0]) == arrow_table_with_null.num_rows - - -def test_bin_pack_record_batches_microbatched(arrow_table_with_null: pa.Table) -> None: - # repeat the per-row batches so we have many small inputs to pack - batches = list(arrow_table_with_null.to_batches(max_chunksize=1)) * 5 - bin_size = arrow_table_with_null.nbytes // 2 # forces multiple bins - bins = list(bin_pack_record_batches(iter(batches), target_file_size=bin_size)) - assert len(bins) > 1 - assert sum(b.num_rows for bin_ in bins for b in bin_) == arrow_table_with_null.num_rows * 5 - # All but the last bin should have crossed the size threshold. - for bin_ in bins[:-1]: - assert sum(b.nbytes for b in bin_) >= bin_size - - -def test_bin_pack_record_batches_empty() -> None: - assert list(bin_pack_record_batches(iter([]), target_file_size=1024)) == [] - - -def test_bin_pack_record_batches_is_lazy(arrow_table_with_null: pa.Table) -> None: - # Streams are single-pass: confirm the helper consumes its input batch-by-batch - # rather than materialising the whole iterator before yielding the first bin. - consumed: list[int] = [] - - def tracking_iter() -> Iterator[pa.RecordBatch]: - for i, batch in enumerate(arrow_table_with_null.to_batches(max_chunksize=1)): - consumed.append(i) - yield batch - - target = max(1, arrow_table_with_null.nbytes // 4) - bins_iter = bin_pack_record_batches(tracking_iter(), target_file_size=target) - first_bin = next(bins_iter) - assert len(first_bin) >= 1 - # Generator should not have walked the entire input upon yielding the first bin - assert len(consumed) < arrow_table_with_null.num_rows - list(bins_iter) - assert len(consumed) == arrow_table_with_null.num_rows - - def test_schema_mismatch_type(table_schema_simple: Schema) -> None: other_schema = pa.schema( (