Enable Avro compression

PyIceberg did not compress the Avro. This will make gzip/deflate
the same as in Java.

Author: Fokko

pyiceberg/avro/codecs/__init__.py

@@ -26,15 +26,19 @@
 
 from __future__ import annotations
 
-from typing import Dict, Optional, Type
+from typing import Dict, Literal, Optional, Type, TypeAlias
 
 from pyiceberg.avro.codecs.bzip2 import BZip2Codec
 from pyiceberg.avro.codecs.codec import Codec
 from pyiceberg.avro.codecs.deflate import DeflateCodec
 from pyiceberg.avro.codecs.snappy_codec import SnappyCodec
 from pyiceberg.avro.codecs.zstandard_codec import ZStandardCodec
 
-KNOWN_CODECS: Dict[str, Optional[Type[Codec]]] = {
+AvroCompressionCodec: TypeAlias = Literal["null", "bzip2", "snappy", "zstandard", "deflate"]
+
+AVRO_CODEC_KEY = "avro.codec"
+
+KNOWN_CODECS: Dict[AvroCompressionCodec, Optional[Type[Codec]]] = {
     "null": None,
     "bzip2": BZip2Codec,
     "snappy": SnappyCodec,

pyiceberg/avro/file.py

@@ -35,7 +35,7 @@
     TypeVar,
 )
 
-from pyiceberg.avro.codecs import KNOWN_CODECS
+from pyiceberg.avro.codecs import AVRO_CODEC_KEY, KNOWN_CODECS
 from pyiceberg.avro.codecs.codec import Codec
 from pyiceberg.avro.decoder import BinaryDecoder, new_decoder
 from pyiceberg.avro.encoder import BinaryEncoder
@@ -69,7 +69,6 @@
     NestedField(field_id=300, name="sync", field_type=FixedType(length=SYNC_SIZE), required=True),
 )
 
-_CODEC_KEY = "avro.codec"
 _SCHEMA_KEY = "avro.schema"
 
 
@@ -92,11 +91,13 @@ def compression_codec(self) -> Optional[Type[Codec]]:
         In the case of a null codec, we return a None indicating that we
         don't need to compress/decompress.
         """
-        codec_name = self.meta.get(_CODEC_KEY, "null")
+        from pyiceberg.table import TableProperties
+
+        codec_name = self.meta.get(AVRO_CODEC_KEY, TableProperties.WRITE_AVRO_COMPRESSION_DEFAULT)
         if codec_name not in KNOWN_CODECS:
             raise ValueError(f"Unsupported codec: {codec_name}")
 
-        return KNOWN_CODECS[codec_name]
+        return KNOWN_CODECS[codec_name]  # type: ignore
 
     def get_schema(self) -> Schema:
         if _SCHEMA_KEY in self.meta:
@@ -276,11 +277,41 @@ def __exit__(
         self.output_stream.close()
 
     def _write_header(self) -> None:
+        from pyiceberg.table import TableProperties
+
+        codec = self.metadata.get(AVRO_CODEC_KEY, TableProperties.WRITE_AVRO_COMPRESSION_DEFAULT)
+        if codec == "gzip":
+            codec = "deflate"
+
         json_schema = json.dumps(AvroSchemaConversion().iceberg_to_avro(self.file_schema, schema_name=self.schema_name))
+<<<<<<< Updated upstream
         meta = {**self.metadata, _SCHEMA_KEY: json_schema, _CODEC_KEY: "null"}
         header = AvroFileHeader(MAGIC, meta, self.sync_bytes)
+=======
+        header = AvroFileHeader(
+            magic=MAGIC, meta={**self.metadata, _SCHEMA_KEY: json_schema, AVRO_CODEC_KEY: codec}, sync=self.sync_bytes
+        )
+>>>>>>> Stashed changes
         construct_writer(META_SCHEMA).write(self.encoder, header)
 
+    def compression_codec(self) -> Optional[Type[Codec]]:
+        """Get the file's compression codec algorithm from the file's metadata.
+
+        In the case of a null codec, we return a None indicating that we
+        don't need to compress/decompress.
+        """
+        from pyiceberg.table import TableProperties
+
+        codec_name = self.metadata.get(AVRO_CODEC_KEY, TableProperties.WRITE_AVRO_COMPRESSION_DEFAULT)
+
+        if codec_name == "gzip":
+            codec_name = "deflate"
+
+        if codec_name not in KNOWN_CODECS:
+            raise ValueError(f"Unsupported codec: {codec_name}")
+
+        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)
@@ -289,6 +320,13 @@ def write_block(self, objects: List[D]) -> None:
         block_content = in_memory.getvalue()
 
         self.encoder.write_int(len(objects))
-        self.encoder.write_int(len(block_content))
-        self.encoder.write(block_content)
+
+        if codec := self.compression_codec():
+            content, content_length = codec.compress(block_content)
+            self.encoder.write_int(content_length)
+            self.encoder.write(content)
+        else:
+            self.encoder.write_int(len(block_content))
+            self.encoder.write(block_content)
+
         self.encoder.write(self.sync_bytes)

pyiceberg/manifest.py

@@ -37,6 +37,7 @@
 from cachetools.keys import hashkey
 from pydantic_core import to_json
 
+from pyiceberg.avro.codecs import AVRO_CODEC_KEY, AvroCompressionCodec
 from pyiceberg.avro.file import AvroFile, AvroOutputFile
 from pyiceberg.conversions import to_bytes
 from pyiceberg.exceptions import ValidationError
@@ -799,8 +800,16 @@ class ManifestWriter(ABC):
     _min_sequence_number: Optional[int]
     _partitions: List[Record]
     _reused_entry_wrapper: ManifestEntry
+    _compression: AvroCompressionCodec
 
-    def __init__(self, spec: PartitionSpec, schema: Schema, output_file: OutputFile, snapshot_id: int) -> None:
+    def __init__(
+        self,
+        spec: PartitionSpec,
+        schema: Schema,
+        output_file: OutputFile,
+        snapshot_id: int,
+        avro_compression: AvroCompressionCodec,
+    ) -> None:
         self.closed = False
         self._spec = spec
         self._schema = schema
@@ -815,6 +824,11 @@ def __init__(self, spec: PartitionSpec, schema: Schema, output_file: OutputFile,
         self._deleted_rows = 0
         self._min_sequence_number = None
         self._partitions = []
+<<<<<<< Updated upstream
+=======
+        self._reused_entry_wrapper = ManifestEntry()
+        self._compression = avro_compression
+>>>>>>> Stashed changes
 
     def __enter__(self) -> ManifestWriter:
         """Open the writer."""
@@ -850,6 +864,7 @@ def _meta(self) -> Dict[str, str]:
             "partition-spec": to_json(self._spec.fields).decode("utf-8"),
             "partition-spec-id": str(self._spec.spec_id),
             "format-version": str(self.version),
+            AVRO_CODEC_KEY: self._compression,
         }
 
     def _with_partition(self, format_version: TableVersion) -> Schema:
@@ -961,13 +976,15 @@ def existing(self, entry: ManifestEntry) -> ManifestWriter:
 
 
 class ManifestWriterV1(ManifestWriter):
-    def __init__(self, spec: PartitionSpec, schema: Schema, output_file: OutputFile, snapshot_id: int):
-        super().__init__(
-            spec,
-            schema,
-            output_file,
-            snapshot_id,
-        )
+    def __init__(
+        self,
+        spec: PartitionSpec,
+        schema: Schema,
+        output_file: OutputFile,
+        snapshot_id: int,
+        avro_compression: AvroCompressionCodec,
+    ):
+        super().__init__(spec, schema, output_file, snapshot_id, avro_compression)
 
     def content(self) -> ManifestContent:
         return ManifestContent.DATA
@@ -981,8 +998,15 @@ def prepare_entry(self, entry: ManifestEntry) -> ManifestEntry:
 
 
 class ManifestWriterV2(ManifestWriter):
-    def __init__(self, spec: PartitionSpec, schema: Schema, output_file: OutputFile, snapshot_id: int):
-        super().__init__(spec, schema, output_file, snapshot_id)
+    def __init__(
+        self,
+        spec: PartitionSpec,
+        schema: Schema,
+        output_file: OutputFile,
+        snapshot_id: int,
+        avro_compression: AvroCompressionCodec,
+    ):
+        super().__init__(spec, schema, output_file, snapshot_id, avro_compression)
 
     def content(self) -> ManifestContent:
         return ManifestContent.DATA
@@ -1008,12 +1032,17 @@ def prepare_entry(self, entry: ManifestEntry) -> ManifestEntry:
 
 
 def write_manifest(
-    format_version: TableVersion, spec: PartitionSpec, schema: Schema, output_file: OutputFile, snapshot_id: int
+    format_version: TableVersion,
+    spec: PartitionSpec,
+    schema: Schema,
+    output_file: OutputFile,
+    snapshot_id: int,
+    avro_compression: AvroCompressionCodec,
 ) -> ManifestWriter:
     if format_version == 1:
-        return ManifestWriterV1(spec, schema, output_file, snapshot_id)
+        return ManifestWriterV1(spec, schema, output_file, snapshot_id, avro_compression)
     elif format_version == 2:
-        return ManifestWriterV2(spec, schema, output_file, snapshot_id)
+        return ManifestWriterV2(spec, schema, output_file, snapshot_id, avro_compression)
     else:
         raise ValueError(f"Cannot write manifest for table version: {format_version}")
 
@@ -1063,14 +1092,21 @@ def add_manifests(self, manifest_files: List[ManifestFile]) -> ManifestListWrite
 
 
 class ManifestListWriterV1(ManifestListWriter):
-    def __init__(self, output_file: OutputFile, snapshot_id: int, parent_snapshot_id: Optional[int]):
+    def __init__(
+        self,
+        output_file: OutputFile,
+        snapshot_id: int,
+        parent_snapshot_id: Optional[int],
+        compression: AvroCompressionCodec,
+    ):
         super().__init__(
             format_version=1,
             output_file=output_file,
             meta={
                 "snapshot-id": str(snapshot_id),
                 "parent-snapshot-id": str(parent_snapshot_id) if parent_snapshot_id is not None else "null",
                 "format-version": "1",
+                AVRO_CODEC_KEY: compression,
             },
         )
 
@@ -1084,7 +1120,14 @@ class ManifestListWriterV2(ManifestListWriter):
     _commit_snapshot_id: int
     _sequence_number: int
 
-    def __init__(self, output_file: OutputFile, snapshot_id: int, parent_snapshot_id: Optional[int], sequence_number: int):
+    def __init__(
+        self,
+        output_file: OutputFile,
+        snapshot_id: int,
+        parent_snapshot_id: Optional[int],
+        sequence_number: int,
+        compression: AvroCompressionCodec,
+    ):
         super().__init__(
             format_version=2,
             output_file=output_file,
@@ -1093,6 +1136,7 @@ def __init__(self, output_file: OutputFile, snapshot_id: int, parent_snapshot_id
                 "parent-snapshot-id": str(parent_snapshot_id) if parent_snapshot_id is not None else "null",
                 "sequence-number": str(sequence_number),
                 "format-version": "2",
+                AVRO_CODEC_KEY: compression,
             },
         )
         self._commit_snapshot_id = snapshot_id
@@ -1127,12 +1171,13 @@ def write_manifest_list(
     snapshot_id: int,
     parent_snapshot_id: Optional[int],
     sequence_number: Optional[int],
+    avro_compression: AvroCompressionCodec,
 ) -> ManifestListWriter:
     if format_version == 1:
-        return ManifestListWriterV1(output_file, snapshot_id, parent_snapshot_id)
+        return ManifestListWriterV1(output_file, snapshot_id, parent_snapshot_id, avro_compression)
     elif format_version == 2:
         if sequence_number is None:
             raise ValueError(f"Sequence-number is required for V2 tables: {sequence_number}")
-        return ManifestListWriterV2(output_file, snapshot_id, parent_snapshot_id, sequence_number)
+        return ManifestListWriterV2(output_file, snapshot_id, parent_snapshot_id, sequence_number, avro_compression)
     else:
         raise ValueError(f"Cannot write manifest list for table version: {format_version}")

pyiceberg/serializers.py

@@ -19,13 +19,15 @@
 import codecs
 import gzip
 from abc import ABC, abstractmethod
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
 
 from pyiceberg.io import InputFile, InputStream, OutputFile
-from pyiceberg.table.metadata import TableMetadata, TableMetadataUtil
 from pyiceberg.typedef import UTF8
 from pyiceberg.utils.config import Config
 
+if TYPE_CHECKING:
+    from pyiceberg.table.metadata import TableMetadata
+
 GZIP = "gzip"
 
 
@@ -79,7 +81,7 @@ class FromByteStream:
     @staticmethod
     def table_metadata(
         byte_stream: InputStream, encoding: str = UTF8, compression: Compressor = NOOP_COMPRESSOR
-    ) -> TableMetadata:
+    ) -> "TableMetadata":
         """Instantiate a TableMetadata object from a byte stream.
 
         Args:
@@ -92,14 +94,16 @@ def table_metadata(
             json_bytes = reader(byte_stream)
             metadata = json_bytes.read()
 
+        from pyiceberg.table.metadata import TableMetadataUtil
+
         return TableMetadataUtil.parse_raw(metadata)
 
 
 class FromInputFile:
     """A collection of methods that deserialize InputFiles into Iceberg objects."""
 
     @staticmethod
-    def table_metadata(input_file: InputFile, encoding: str = UTF8) -> TableMetadata:
+    def table_metadata(input_file: InputFile, encoding: str = UTF8) -> "TableMetadata":
         """Create a TableMetadata instance from an input file.
 
         Args:
@@ -120,7 +124,7 @@ class ToOutputFile:
     """A collection of methods that serialize Iceberg objects into files given an OutputFile instance."""
 
     @staticmethod
-    def table_metadata(metadata: TableMetadata, output_file: OutputFile, overwrite: bool = False) -> None:
+    def table_metadata(metadata: "TableMetadata", output_file: OutputFile, overwrite: bool = False) -> None:
         """Write a TableMetadata instance to an output file.
 
         Args:

pyiceberg/table/__init__.py

@@ -192,6 +192,9 @@ class TableProperties:
     WRITE_TARGET_FILE_SIZE_BYTES = "write.target-file-size-bytes"
     WRITE_TARGET_FILE_SIZE_BYTES_DEFAULT = 512 * 1024 * 1024  # 512 MB
 
+    WRITE_AVRO_COMPRESSION = "write.avro.compression-codec"
+    WRITE_AVRO_COMPRESSION_DEFAULT = "gzip"
+
     DEFAULT_WRITE_METRICS_MODE = "write.metadata.metrics.default"
     DEFAULT_WRITE_METRICS_MODE_DEFAULT = "truncate(16)"