abnobdoss commented on code in PR #3303:
URL: https://github.com/apache/iceberg-python/pull/3303#discussion_r3331231702
##########
tests/avro/test_encoder.py:
##########
@@ -78,57 +82,64 @@ def test_write_int() -> None:
assert buffer[28:36] == b"\xc4\xa0\xb2\xae\x83\xf8\xe4\x7c"
-def test_write_float() -> None:
- output = io.BytesIO()
- encoder = BinaryEncoder(output)
-
+def test_write_float(encoder: Any) -> None:
_input = 3.14159265359
-
encoder.write_float(_input)
-
- assert output.getbuffer() == struct.pack("<f", _input)
+ assert encoder.getvalue() == struct.pack("<f", _input)
-def test_write_double() -> None:
- output = io.BytesIO()
- encoder = BinaryEncoder(output)
-
+def test_write_double(encoder: Any) -> None:
_input = 3.14159265359
-
encoder.write_double(_input)
-
- assert output.getbuffer() == struct.pack("<d", _input)
+ assert encoder.getvalue() == struct.pack("<d", _input)
-def test_write_bytes() -> None:
- output = io.BytesIO()
- encoder = BinaryEncoder(output)
-
+def test_write_bytes(encoder: Any) -> None:
_input = b"\x12\x34\x56"
-
encoder.write_bytes(_input)
+ assert encoder.getvalue() == b"".join([b"\x06", _input])
- assert output.getbuffer() == b"".join([b"\x06", _input])
-
-
-def test_write_utf8() -> None:
- output = io.BytesIO()
- encoder = BinaryEncoder(output)
+def test_write_utf8(encoder: Any) -> None:
_input = "That, my liege, is how we know the Earth to be banana-shaped."
bin_input = _input.encode()
encoder.write_utf8(_input)
-
- assert output.getbuffer() == b"".join([b"\x7a", bin_input])
+ assert encoder.getvalue() == b"".join([b"\x7a", bin_input])
-def test_write_uuid() -> None:
- output = io.BytesIO()
- encoder = BinaryEncoder(output)
-
+def test_write_uuid(encoder: Any) -> None:
_input = uuid.UUID("12345678-1234-5678-1234-567812345678")
encoder.write_uuid(_input)
-
- buf = output.getbuffer()
+ buf = encoder.getvalue()
assert len(buf) == 16
- assert buf.tobytes() == b"\x124Vx\x124Vx\x124Vx\x124Vx"
+ assert buf == b"\x124Vx\x124Vx\x124Vx\x124Vx"
+
+
[email protected](
+ "v",
+ [0, 1, -1, 63, 64, -64, -65, 127, 128, -128, 2**31 - 1, -(2**31), 2**62,
-(2**62), 2**63 - 1, -(2**63)],
+)
+def test_int_round_trip(encoder: Any, v: int) -> None:
Review Comment:
Nice! This is actually stricter than the python encoder method: python
currently accepts `2**63` and `-(2**63) - 1` even though those aren’t valid
Avro longs. Should we align them / add out-of-range tests so behavior doesn’t
depend on which encoder is used?
##########
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
+
+
[email protected]
+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"))
Review Comment:
minor: should this use UTF8 from pyiceberg.typedef?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
