perf(avro): buffer records into size-bounded blocks

Author: mbutrovich

pyiceberg/avro/file.py

@@ -54,6 +54,9 @@
 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 @@ def __init__(
         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 @@ def __init__(
             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 @@ def __enter__(self) -> AvroOutputFile[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 @@ def compression_codec(self) -> type[Codec] | None:
         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 @@ def write_block(self, objects: list[D]) -> None:
 
         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()

tests/avro/test_file.py

@@ -23,11 +23,11 @@
 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 field_uuid(self) -> UUID:
     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,
+        ),
+    )
+
+
+@pytest.mark.parametrize("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)) == []
+
+
+@pytest.mark.parametrize("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

tests/utils/test_manifest.py

@@ -25,7 +25,7 @@
 
 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"
 
 
+@pytest.mark.parametrize("format_version", [1, 2])
+@pytest.mark.parametrize("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)]
+
+
+@pytest.mark.parametrize("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