Add mapping

Author: Fokko

pyiceberg/avro/codecs/__init__.py

@@ -47,3 +47,6 @@
     "zstandard": ZStandardCodec,
     "deflate": DeflateCodec,
 }
+
+# Map to convert the naming from Iceberg to Avro
+CODEC_MAPPING_ICEBERG_TO_AVRO: Dict[str, str] = {"gzip": "deflate", "zstd": "zstandard"}

pyiceberg/avro/file.py

@@ -35,7 +35,7 @@
     TypeVar,
 )
 
-from pyiceberg.avro.codecs import AVRO_CODEC_KEY, KNOWN_CODECS
+from pyiceberg.avro.codecs import AVRO_CODEC_KEY, CODEC_MAPPING_ICEBERG_TO_AVRO, KNOWN_CODECS
 from pyiceberg.avro.codecs.codec import Codec
 from pyiceberg.avro.decoder import BinaryDecoder, new_decoder
 from pyiceberg.avro.encoder import BinaryEncoder
@@ -279,13 +279,13 @@ def __exit__(
     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"
+        codec_name = self.metadata.get(AVRO_CODEC_KEY, TableProperties.WRITE_AVRO_COMPRESSION_DEFAULT)
+        if avro_codec_name := CODEC_MAPPING_ICEBERG_TO_AVRO.get(codec_name):
+            codec_name = avro_codec_name
 
         json_schema = json.dumps(AvroSchemaConversion().iceberg_to_avro(self.file_schema, schema_name=self.schema_name))
 
-        meta = {**self.metadata, _SCHEMA_KEY: json_schema, AVRO_CODEC_KEY: codec}
+        meta = {**self.metadata, _SCHEMA_KEY: json_schema, AVRO_CODEC_KEY: codec_name}
         header = AvroFileHeader(MAGIC, meta, self.sync_bytes)
         construct_writer(META_SCHEMA).write(self.encoder, header)
 
@@ -299,8 +299,8 @@ def compression_codec(self) -> Optional[Type[Codec]]:
 
         codec_name = self.metadata.get(AVRO_CODEC_KEY, TableProperties.WRITE_AVRO_COMPRESSION_DEFAULT)
 
-        if codec_name == "gzip":
-            codec_name = "deflate"
+        if avro_codec_name := CODEC_MAPPING_ICEBERG_TO_AVRO.get(codec_name):
+            codec_name = avro_codec_name
 
         if codec_name not in KNOWN_CODECS:
             raise ValueError(f"Unsupported codec: {codec_name}")