Add Cython BinaryEncoder for Avro block encoding

.gitignore

@@ -45,5 +45,6 @@ htmlcov
 .ipynb_checkpoints/
 
 pyiceberg/avro/decoder_fast.c
+pyiceberg/avro/encoder_fast.c
 pyiceberg/avro/*.html
 pyiceberg/avro/*.so

MANIFEST.in

@@ -20,6 +20,7 @@ recursive-include pyiceberg *.pyx *.c
 
 # Exclude generated Cython C file
 exclude pyiceberg/avro/decoder_fast.c
+exclude pyiceberg/avro/encoder_fast.c
 
 # Include test files in sdist
 recursive-include tests *

pyiceberg/avro/encoder.py

@@ -78,3 +78,29 @@ def write_uuid(self, uuid: UUID) -> None:
 
     def write_unknown(self, _: Any) -> None:
         """Nulls are written as 0 bytes in avro, so we do nothing."""
+
+
+class MemoryBinaryEncoder(BinaryEncoder):
+    """BinaryEncoder that writes to an owned in-memory buffer."""
+
+    def __init__(self) -> None:
+        import io
+
+        self._buffer = io.BytesIO()
+        super().__init__(self._buffer)
+
+    def getvalue(self) -> bytes:
+        return self._buffer.getvalue()
+
+
+def new_memory_encoder() -> "CythonBinaryEncoder | MemoryBinaryEncoder":  # type: ignore[name-defined]  # noqa: F821
+    try:
+        from pyiceberg.avro.encoder_fast import CythonBinaryEncoder
+
+        return CythonBinaryEncoder()
+    except ModuleNotFoundError:
+        import warnings
+
+        warnings.warn("Falling back to pure Python Avro encoder, missing Cython implementation", stacklevel=2)
+
+        return MemoryBinaryEncoder()

pyiceberg/avro/encoder_fast.pyi

@@ -0,0 +1,31 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from typing import Any
+from uuid import UUID
+
+class CythonBinaryEncoder:
+    def __init__(self) -> None: ...
+    def getvalue(self) -> bytes: ...
+    def write(self, b: bytes) -> None: ...
+    def write_boolean(self, v: bool) -> None: ...
+    def write_int(self, v: int) -> None: ...
+    def write_float(self, v: float) -> None: ...
+    def write_double(self, v: float) -> None: ...
+    def write_bytes(self, b: bytes) -> None: ...
+    def write_utf8(self, s: str) -> None: ...
+    def write_uuid(self, uuid: UUID) -> None: ...
+    def write_unknown(self, _: Any) -> None: ...

pyiceberg/avro/encoder_fast.pyx

@@ -0,0 +1,119 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import cython
+from cpython.bytes cimport PyBytes_FromStringAndSize, PyBytes_AS_STRING, PyBytes_GET_SIZE
+from cpython.mem cimport PyMem_Malloc, PyMem_Realloc, PyMem_Free
+from libc.string cimport memcpy
+from libc.stdint cimport uint8_t, uint64_t, int64_t
+
+from pyiceberg.avro import STRUCT_DOUBLE, STRUCT_FLOAT
+
+cdef uint64_t _INITIAL_CAPACITY = 1024
+
+
+@cython.final
+cdef class CythonBinaryEncoder:
+    """In-memory BinaryEncoder that writes to a growable C buffer.
+
+    Drop-in replacement for BinaryEncoder for the block-encoding path:
+    exposes the same write_* methods the Writer tree calls, plus
+    getvalue() to materialise the encoded bytes once at the end.
+    """
+
+    cdef unsigned char *_data
+    cdef uint64_t _size
+    cdef uint64_t _capacity
+
+    def __cinit__(self):
+        self._data = <unsigned char *> PyMem_Malloc(_INITIAL_CAPACITY)
+        if not self._data:
+            raise MemoryError()
+        self._size = 0
+        self._capacity = _INITIAL_CAPACITY
+
+    def __dealloc__(self):
+        PyMem_Free(self._data)
+
+    cdef inline int _ensure(self, uint64_t n) except -1:
+        cdef uint64_t need = self._size + n
+        if need <= self._capacity:
+            return 0
+        cdef uint64_t cap = self._capacity
+        while cap < need:
+            cap <<= 1
+        cdef unsigned char *grown = <unsigned char *> PyMem_Realloc(self._data, cap)
+        if not grown:
+            raise MemoryError()
+        self._data = grown
+        self._capacity = cap
+        return 0
+
+    cpdef bytes getvalue(self):
+        return PyBytes_FromStringAndSize(<char *> self._data, self._size)
+
+    cpdef void write(self, bytes b):
+        cdef Py_ssize_t n = PyBytes_GET_SIZE(b)
+        self._ensure(n)
+        memcpy(self._data + self._size, PyBytes_AS_STRING(b), n)
+        self._size += n
+
+    cpdef void write_boolean(self, bint v):
+        self._ensure(1)
+        self._data[self._size] = 1 if v else 0
+        self._size += 1
+
+    cpdef void write_int(self, int64_t v):
+        # zigzag then base-128 varint; a 64-bit value needs at most 10 bytes
+        self._ensure(10)
+        cdef uint64_t uv = <uint64_t> v
+        cdef uint64_t datum = (uv << 1) ^ (0 - (uv >> 63))
+        cdef unsigned char *p = self._data + self._size
+        while datum & <uint64_t> ~0x7F:
+            p[0] = <uint8_t> ((datum & 0x7F) | 0x80)
+            p += 1
+            datum >>= 7
+        p[0] = <uint8_t> datum
+        p += 1
+        self._size = p - self._data
+
+    def write_float(self, v: float) -> None:
+        self.write(STRUCT_FLOAT.pack(v))
+
+    def write_double(self, v: float) -> None:
+        self.write(STRUCT_DOUBLE.pack(v))
+
+    cpdef void write_bytes(self, b):
+        cdef bytes bb = bytes(b) if type(b) is not bytes else b
+        cdef Py_ssize_t n = PyBytes_GET_SIZE(bb)
+        self.write_int(n)
+        self._ensure(n)
+        memcpy(self._data + self._size, PyBytes_AS_STRING(bb), n)
+        self._size += n
+
+    def write_utf8(self, s) -> None:
+        self.write_bytes(s.encode("utf-8"))
+
+    def write_uuid(self, uuid) -> None:
+        cdef bytes b = uuid.bytes
+        if PyBytes_GET_SIZE(b) != 16:
+            raise ValueError(f"Expected UUID to have 16 bytes, got: len({b!r})")
+        self._ensure(16)
+        memcpy(self._data + self._size, PyBytes_AS_STRING(b), 16)
+        self._size += 16
+
+    def write_unknown(self, _) -> None:
+        pass

pyiceberg/avro/file.py

@@ -19,7 +19,6 @@
 
 from __future__ import annotations
 
-import io
 import json
 import os
 from collections.abc import Callable
@@ -34,7 +33,7 @@
 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
+from pyiceberg.avro.encoder import BinaryEncoder, new_memory_encoder
 from pyiceberg.avro.reader import Reader
 from pyiceberg.avro.resolver import construct_reader, construct_writer, resolve_reader, resolve_writer
 from pyiceberg.avro.writer import Writer
@@ -300,11 +299,10 @@ 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)
+        block_content_encoder = new_memory_encoder()
         for obj in objects:
             self.writer.write(block_content_encoder, obj)
-        block_content = in_memory.getvalue()
+        block_content = block_content_encoder.getvalue()
 
         self.encoder.write_int(len(objects))
 

setup.py

@@ -59,7 +59,13 @@ def make_release_tree(self, base_dir: str, files: list[str]) -> None:
             [os.path.join(package_path, "avro", "decoder_fast.pyx")],
             extra_compile_args=extra_compile_args,
             language="c",
-        )
+        ),
+        Extension(
+            "pyiceberg.avro.encoder_fast",
+            [os.path.join(package_path, "avro", "encoder_fast.pyx")],
+            extra_compile_args=extra_compile_args,
+            language="c",
+        ),
     ]
 
     ext_modules = cythonize(