This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-4008-3ec3b3f50de210f0dad31f18a627ad530fd744e3 in repository https://gitbox.apache.org/repos/asf/iceberg-python.git
commit 80e9041c6a4a3ae27b8bbb7b6a7ab0295766d2a4 Author: Matt Butrovich <[email protected]> AuthorDate: Fri Oct 2 21:14:48 2026 +0000 perf(avro): buffer records into size-bounded blocks (#4008) --- pyiceberg/avro/file.py | 32 +++++++++++--- tests/avro/test_file.py | 95 ++++++++++++++++++++++++++++++++++++++- tests/utils/test_manifest.py | 103 ++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 221 insertions(+), 9 deletions(-) diff --git a/pyiceberg/avro/file.py b/pyiceberg/avro/file.py index 7db92818f..3184e3950 100644 --- a/pyiceberg/avro/file.py +++ b/pyiceberg/avro/file.py @@ -54,6 +54,9 @@ VERSION = 1 MAGIC = bytes(b"Obj" + bytearray([VERSION])) MAGIC_SIZE = len(MAGIC) SYNC_SIZE = 16 +# Approximate uncompressed bytes buffered per block, matching the default sync interval of the Avro +# Java writer that Iceberg Java uses (DataFileConstants.DEFAULT_SYNC_INTERVAL). +DEFAULT_SYNC_INTERVAL = 4000 * SYNC_SIZE META_SCHEMA = StructType( NestedField(name="magic", field_id=100, field_type=FixedType(length=MAGIC_SIZE), required=True), NestedField( @@ -238,6 +241,7 @@ class AvroOutputFile(Generic[D]): schema_name: str, record_schema: Schema | None = None, metadata: dict[str, str] = EMPTY_DICT, + sync_interval: int = DEFAULT_SYNC_INTERVAL, ) -> None: self.output_file = output_file self.file_schema = file_schema @@ -249,6 +253,12 @@ class AvroOutputFile(Generic[D]): else resolve_writer(record_schema=record_schema, file_schema=self.file_schema) ) self.metadata = metadata + if sync_interval <= 0: + raise ValueError(f"Sync interval must be positive: {sync_interval}") + self.sync_interval = sync_interval + self._block = io.BytesIO() + self._block_encoder = BinaryEncoder(output_stream=self._block) + self._block_records = 0 def __enter__(self) -> AvroOutputFile[D]: """ @@ -266,6 +276,7 @@ class AvroOutputFile(Generic[D]): def __exit__(self, exctype: type[BaseException] | None, excinst: BaseException | None, exctb: TracebackType | None) -> None: """Perform cleanup when exiting the scope of a 'with' statement.""" + self._flush_block() self.output_stream.close() def _write_header(self) -> None: @@ -300,13 +311,18 @@ class AvroOutputFile(Generic[D]): return KNOWN_CODECS[codec_name] # type: ignore def write_block(self, objects: list[D]) -> None: - in_memory = io.BytesIO() - block_content_encoder = BinaryEncoder(output_stream=in_memory) for obj in objects: - self.writer.write(block_content_encoder, obj) - block_content = in_memory.getvalue() + self.writer.write(self._block_encoder, obj) + self._block_records += 1 + if self._block.tell() >= self.sync_interval: + self._flush_block() - self.encoder.write_int(len(objects)) + def _flush_block(self) -> None: + if self._block_records == 0: + return + + block_content = self._block.getvalue() + self.encoder.write_int(self._block_records) if codec := self.compression_codec(): content, content_length = codec.compress(block_content) @@ -318,5 +334,9 @@ class AvroOutputFile(Generic[D]): self.encoder.write(self.sync_bytes) + self._block.seek(0) + self._block.truncate() + self._block_records = 0 + def tell(self) -> int: - return self.output_stream.tell() + return self.output_stream.tell() + self._block.tell() diff --git a/tests/avro/test_file.py b/tests/avro/test_file.py index 2d3ddeefa..175089c30 100644 --- a/tests/avro/test_file.py +++ b/tests/avro/test_file.py @@ -23,11 +23,11 @@ from typing import Any from uuid import UUID import pytest -from fastavro import reader, writer +from fastavro import block_reader, reader, writer import pyiceberg.avro.file as avro from pyiceberg.avro.codecs.deflate import DeflateCodec -from pyiceberg.avro.file import AvroFileHeader +from pyiceberg.avro.file import DEFAULT_SYNC_INTERVAL, AvroFileHeader from pyiceberg.io.pyarrow import PyArrowFileIO from pyiceberg.manifest import ( DEFAULT_BLOCK_SIZE, @@ -453,3 +453,94 @@ def test_all_primitive_types(is_required: bool) -> None: for idx, field in enumerate(all_primitives_schema.as_struct()): assert record[idx] == avro_entry[idx], f"Invalid {field}" assert record[idx] == avro_entry_read_with_fastavro[idx], f"Invalid {field} read with fastavro" + + +def manifest_entry(index: int) -> ManifestEntry: + return ManifestEntry.from_args( + status=ManifestEntryStatus.ADDED, + snapshot_id=8638475580105682862, + sequence_number=0, + file_sequence_number=0, + data_file=DataFile.from_args( + content=DataFileContent.DATA, + file_path=f"s3://some-path/some-file-{index}.parquet", + file_format=FileFormat.PARQUET, + partition=Record(), + record_count=131327, + file_size_in_bytes=220669226, + ), + ) + + [email protected]("sync_interval", [0, -1]) +def test_avro_output_file_rejects_non_positive_sync_interval(sync_interval: int) -> None: + with TemporaryDirectory() as tmpdir: + with pytest.raises(ValueError, match="Sync interval must be positive"): + avro.AvroOutputFile[ManifestEntry]( + output_file=PyArrowFileIO().new_output(tmpdir + "/manifest_entry.avro"), + file_schema=MANIFEST_ENTRY_SCHEMAS[2], + schema_name="manifest_entry", + sync_interval=sync_interval, + ) + + +def test_avro_output_file_writes_no_block_without_records() -> None: + with TemporaryDirectory() as tmpdir: + tmp_avro_file = tmpdir + "/manifest_entry.avro" + + with avro.AvroOutputFile[ManifestEntry]( + output_file=PyArrowFileIO().new_output(tmp_avro_file), + file_schema=MANIFEST_ENTRY_SCHEMAS[2], + schema_name="manifest_entry", + ): + pass + + with open(tmp_avro_file, "rb") as fo: + assert list(block_reader(fo)) == [] + + [email protected]("sync_interval", [1, DEFAULT_SYNC_INTERVAL]) +def test_avro_output_file_round_trips_across_blocks(sync_interval: int) -> None: + entries = [manifest_entry(index) for index in range(10)] + + with TemporaryDirectory() as tmpdir: + tmp_avro_file = tmpdir + "/manifest_entry.avro" + + with avro.AvroOutputFile[ManifestEntry]( + output_file=PyArrowFileIO().new_output(tmp_avro_file), + file_schema=MANIFEST_ENTRY_SCHEMAS[2], + schema_name="manifest_entry", + sync_interval=sync_interval, + ) as out: + for entry in entries: + out.write_block([entry]) + + with open(tmp_avro_file, "rb") as fo: + blocks = [block.num_records for block in block_reader(fo)] + + with avro.AvroFile[ManifestEntry]( + input_file=PyArrowFileIO().new_input(tmp_avro_file), + read_schema=MANIFEST_ENTRY_SCHEMAS[2], + read_types={-1: ManifestEntry, 2: DataFile}, + ) as avro_reader: + read_entries = list(avro_reader) + + if sync_interval == 1: + # A record that on its own reaches the sync interval is flushed as its own block + assert blocks == [1] * len(entries) + else: + assert blocks == [len(entries)] + assert [entry.data_file.file_path for entry in read_entries] == [entry.data_file.file_path for entry in entries] + + +def test_avro_output_file_tell_includes_buffered_records() -> None: + with TemporaryDirectory() as tmpdir: + with avro.AvroOutputFile[ManifestEntry]( + output_file=PyArrowFileIO().new_output(tmpdir + "/manifest_entry.avro"), + file_schema=MANIFEST_ENTRY_SCHEMAS[2], + schema_name="manifest_entry", + ) as out: + after_header = out.tell() + out.write_block([manifest_entry(0)]) + + assert out.tell() > after_header diff --git a/tests/utils/test_manifest.py b/tests/utils/test_manifest.py index 331146346..45e4e894e 100644 --- a/tests/utils/test_manifest.py +++ b/tests/utils/test_manifest.py @@ -25,7 +25,7 @@ import pytest import pyiceberg.manifest as manifest_module from pyiceberg.avro.codecs import AvroCompressionCodec -from pyiceberg.avro.file import AvroFile, AvroOutputFile +from pyiceberg.avro.file import DEFAULT_SYNC_INTERVAL, AvroFile, AvroOutputFile from pyiceberg.io import load_file_io from pyiceberg.io.pyarrow import PyArrowFileIO from pyiceberg.manifest import ( @@ -1126,6 +1126,107 @@ def test_manifest_writer_tell(format_version: TableVersion) -> None: assert after_entry_bytes > initial_bytes, "Bytes should increase after adding entry" [email protected]("format_version", [1, 2]) [email protected]("compression", ["null", "deflate"]) +def test_write_manifest_fills_blocks(format_version: TableVersion, compression: AvroCompressionCodec) -> None: + io = load_file_io() + test_schema = Schema(NestedField(1, "foo", IntegerType(), False)) + entry_count = 2000 + + with TemporaryDirectory() as tmpdir: + manifest_path = f"{tmpdir}/test-manifest.avro" + with write_manifest( + format_version=format_version, + spec=UNPARTITIONED_PARTITION_SPEC, + schema=test_schema, + output_file=io.new_output(manifest_path), + snapshot_id=1, + avro_compression=compression, + ) as writer: + for i in range(entry_count): + writer.add_entry( + ManifestEntry.from_args( + status=ManifestEntryStatus.ADDED, + snapshot_id=1, + data_file=DataFile.from_args( + content=DataFileContent.DATA, + file_path=f"{tmpdir}/data-{i}.parquet", + file_format=FileFormat.PARQUET, + partition=Record(), + record_count=100, + file_size_in_bytes=1000, + ), + ) + ) + + with open(manifest_path, "rb") as f: + blocks = [(block.num_records, block.size) for block in fastavro.block_reader(f)] + + manifest = ManifestFile.from_args( + manifest_path=manifest_path, + manifest_length=0, + partition_spec_id=0, + added_snapshot_id=1, + sequence_number=0, + partitions=[], + ) + read_entries = manifest.fetch_manifest_entry(io) + + assert sum(num_records for num_records, _ in blocks) == entry_count + assert len(blocks) < entry_count + if compression == "null": + # Only an uncompressed block reports the buffered size that triggers a flush + assert all(size >= DEFAULT_SYNC_INTERVAL for _, size in blocks[:-1]) + assert [entry.data_file.file_path for entry in read_entries] == [f"{tmpdir}/data-{i}.parquet" for i in range(entry_count)] + + [email protected]("format_version", [1, 2]) +def test_write_manifest_list_fills_blocks(format_version: TableVersion) -> None: + io = load_file_io() + manifest_count = 2000 + + with TemporaryDirectory() as tmpdir: + manifest_list_path = f"{tmpdir}/manifest-list.avro" + manifests = [ + ManifestFile.from_args( + manifest_path=f"{tmpdir}/manifest-{i}.avro", + manifest_length=100, + partition_spec_id=0, + content=ManifestContent.DATA, + sequence_number=0, + min_sequence_number=0, + added_snapshot_id=1, + added_files_count=1, + existing_files_count=0, + deleted_files_count=0, + added_rows_count=1, + existing_rows_count=0, + deleted_rows_count=0, + partitions=[], + ) + for i in range(manifest_count) + ] + with write_manifest_list( + format_version=format_version, + output_file=io.new_output(manifest_list_path), + snapshot_id=1, + parent_snapshot_id=None, + sequence_number=0, + avro_compression="null", + ) as writer: + writer.add_manifests(manifests) + + with open(manifest_list_path, "rb") as f: + blocks = [(block.num_records, block.size) for block in fastavro.block_reader(f)] + + read_manifests = list(read_manifest_list(io.new_input(manifest_list_path))) + + assert sum(num_records for num_records, _ in blocks) == manifest_count + assert all(size >= DEFAULT_SYNC_INTERVAL for _, size in blocks[:-1]) + assert len(blocks) > 1 + assert [manifest.manifest_path for manifest in read_manifests] == [manifest.manifest_path for manifest in manifests] + + @pytest.mark.parametrize("format_version", [1, 2]) def test_write_manifest_min_sequence_number_zero(format_version: TableVersion) -> None: # A data sequence number of 0 is a legitimate min for a live file (e.g. files from a
