Python: Add more tests for the Avro writer

python/pyiceberg/avro/decoder.py

@@ -19,7 +19,6 @@
 from datetime import datetime, time
 from io import SEEK_CUR
 from typing import Dict, List
-from uuid import UUID
 
 from pyiceberg.avro import STRUCT_DOUBLE, STRUCT_FLOAT
 from pyiceberg.io import InputStream
@@ -138,10 +137,6 @@ def read_utf8(self) -> str:
         """
         return self.read_bytes().decode("utf-8")
 
-    def read_uuid_from_fixed(self) -> UUID:
-        """Reads a UUID as a fixed[16]."""
-        return UUID(bytes=self.read(16))
-
     def read_time_millis(self) -> time:
         """Reads a milliseconds granularity time from the stream.
 

python/pyiceberg/avro/encoder.py

@@ -15,12 +15,13 @@
 # specific language governing permissions and limitations
 # under the License.
 import decimal
-import struct
-from datetime import date, datetime, time
+from datetime import time
+from uuid import UUID
 
 from pyiceberg.avro import STRUCT_DOUBLE, STRUCT_FLOAT
 from pyiceberg.io import OutputStream
-from pyiceberg.utils.datetime import date_to_days, datetime_to_micros, time_object_to_micros
+from pyiceberg.utils.datetime import time_to_micros
+from pyiceberg.utils.decimal import decimal_to_unscaled
 
 
 class BinaryEncoder:
@@ -64,112 +65,40 @@ def write_decimal_bytes(self, datum: decimal.Decimal) -> None:
 
         Since size of packed value in bytes for signed long is 8, 8 bytes are written.
         """
-        sign, digits, _ = datum.as_tuple()
-
-        unscaled_datum = 0
-        for digit in digits:
-            unscaled_datum = (unscaled_datum * 10) + digit
-
-        bits_req = unscaled_datum.bit_length() + 1
-        if sign:
-            unscaled_datum = (1 << bits_req) - unscaled_datum
-
-        bytes_req = bits_req // 8
-        padding_bits = ~((1 << bits_req) - 1) if sign else 0
-        packed_bits = padding_bits | unscaled_datum
-
-        bytes_req += 1 if (bytes_req << 3) < bits_req else 0
-        self.write_int(bytes_req)
-        for index in range(bytes_req - 1, -1, -1):
-            bits_to_write = packed_bits >> (8 * index)
-            self.write(bytearray([bits_to_write & 0xFF]))
+        unscaled_datum = decimal_to_unscaled(datum)
+        size = (unscaled_datum.bit_length() + 7) // 8
+        bytes_datum = unscaled_datum.to_bytes(length=size, byteorder="big", signed=True)
+        self.write_bytes(bytes_datum)
 
     def write_decimal_fixed(self, datum: decimal.Decimal, size: int) -> None:
         """Decimal in fixed are encoded as size of fixed bytes."""
-        sign, digits, _ = datum.as_tuple()
-
-        unscaled_datum = 0
-        for digit in digits:
-            unscaled_datum = (unscaled_datum * 10) + digit
-
-        bits_req = unscaled_datum.bit_length() + 1
-        size_in_bits = size * 8
-        offset_bits = size_in_bits - bits_req
-
-        mask = 2**size_in_bits - 1
-        bit = 1
-        for _ in range(bits_req):
-            mask ^= bit
-            bit <<= 1
-
-        if bits_req < 8:
-            bytes_req = 1
-        else:
-            bytes_req = bits_req // 8
-            if bits_req % 8 != 0:
-                bytes_req += 1
-        if sign:
-            unscaled_datum = (1 << bits_req) - unscaled_datum
-            unscaled_datum = mask | unscaled_datum
-            for index in range(size - 1, -1, -1):
-                bits_to_write = unscaled_datum >> (8 * index)
-                self.write(bytearray([bits_to_write & 0xFF]))
-        else:
-            for _ in range(offset_bits // 8):
-                self.write(b"\x00")
-            for index in range(bytes_req - 1, -1, -1):
-                bits_to_write = unscaled_datum >> (8 * index)
-                self.write(bytearray([bits_to_write & 0xFF]))
+        unscaled_datum = decimal_to_unscaled(datum)
+        bytes_datum = unscaled_datum.to_bytes(length=size, byteorder="big", signed=True)
+        self.write(bytes_datum)
 
     def write_bytes(self, b: bytes) -> None:
         """Bytes are encoded as a long followed by that many bytes of data."""
         self.write_int(len(b))
-        self.write(struct.pack(f"{len(b)}s", b))
-
-    def write_bytes_fixed(self, b: bytes) -> None:
-        """Writes fixed number of bytes."""
-        self.write(struct.pack(f"{len(b)}s", b))
+        self.write(b)
 
     def write_utf8(self, s: str) -> None:
         """A string is encoded as a long followed by that many bytes of UTF-8 encoded character data."""
         self.write_bytes(s.encode("utf-8"))
 
-    def write_date_int(self, d: date) -> None:
-        """
-        Encode python date object as int.
-
-        It stores the number of days from the unix epoch, 1 January 1970 (ISO calendar).
-        """
-        self.write_int(date_to_days(d))
-
-    def write_time_millis_int(self, dt: time) -> None:
-        """
-        Encode python time object as int.
+    def write_uuid(self, uuid: UUID) -> None:
+        """Write UUID as a fixed[16].
 
-        It stores the number of milliseconds from midnight, 00:00:00.000
+        The uuid logical type represents a random generated universally unique identifier (UUID).
+        An uuid logical type annotates an Avro string. The string has to conform with RFC-4122.
         """
-        self.write_int(int(time_object_to_micros(dt) / 1000))
+        if len(uuid.bytes) != 16:
+            raise ValueError(f"Expected UUID to have 16 bytes, got: len({uuid.bytes!r})")
+        return self.write(uuid.bytes)
 
-    def write_time_micros_long(self, dt: time) -> None:
+    def write_time_micros(self, dt: time) -> None:
         """
         Encode python time object as long.
 
         It stores the number of microseconds from midnight, 00:00:00.000000
         """
-        self.write_int(time_object_to_micros(dt))
-
-    def write_timestamp_millis_long(self, dt: datetime) -> None:
-        """
-        Encode python datetime object as long.
-
-        It stores the number of milliseconds from midnight of unix epoch, 1 January 1970.
-        """
-        self.write_int(int(datetime_to_micros(dt) / 1000))
-
-    def write_timestamp_micros_long(self, dt: datetime) -> None:
-        """
-        Encode python datetime object as long.
-
-        It stores the number of microseconds from midnight of unix epoch, 1 January 1970.
-        """
-        self.write_int(datetime_to_micros(dt))
+        self.write_int(time_to_micros(dt))

python/pyiceberg/avro/file.py

@@ -275,4 +275,4 @@ def write_block(self, objects: List[D]) -> None:
         self.encoder.write_int(len(objects))
         self.encoder.write_int(len(block_content))
         self.encoder.write(block_content)
-        self.encoder.write_bytes_fixed(self.sync_bytes)
+        self.encoder.write(self.sync_bytes)

python/pyiceberg/avro/reader.py

@@ -28,7 +28,7 @@
 from abc import abstractmethod
 from dataclasses import dataclass
 from dataclasses import field as dataclassfield
-from datetime import datetime, time
+from datetime import date, datetime, time
 from decimal import Decimal
 from typing import (
     Any,
@@ -43,6 +43,7 @@
 from pyiceberg.avro.decoder import BinaryDecoder
 from pyiceberg.typedef import StructProtocol
 from pyiceberg.types import StructType
+from pyiceberg.utils.datetime import days_to_date
 from pyiceberg.utils.singleton import Singleton
 
 
@@ -153,8 +154,8 @@ def skip(self, decoder: BinaryDecoder) -> None:
 
 
 class DateReader(Reader):
-    def read(self, decoder: BinaryDecoder) -> int:
-        return decoder.read_int()
+    def read(self, decoder: BinaryDecoder) -> date:
+        return days_to_date(decoder.read_int())
 
     def skip(self, decoder: BinaryDecoder) -> None:
         decoder.skip_int()
@@ -194,7 +195,7 @@ def skip(self, decoder: BinaryDecoder) -> None:
 
 class UUIDReader(Reader):
     def read(self, decoder: BinaryDecoder) -> UUID:
-        return decoder.read_uuid_from_fixed()
+        return UUID(bytes=decoder.read(16))
 
     def skip(self, decoder: BinaryDecoder) -> None:
         decoder.skip(16)

python/pyiceberg/avro/writer.py

@@ -36,6 +36,7 @@
 
 from pyiceberg.avro.encoder import BinaryEncoder
 from pyiceberg.types import StructType
+from pyiceberg.utils.datetime import date_to_days, datetime_to_micros
 from pyiceberg.utils.singleton import Singleton
 
 
@@ -78,22 +79,25 @@ def write(self, encoder: BinaryEncoder, val: float) -> None:
 
 class DateWriter(Writer):
     def write(self, encoder: BinaryEncoder, val: Any) -> None:
-        encoder.write_date_int(val)
+        encoder.write_int(date_to_days(val))
 
 
 class TimeWriter(Writer):
     def write(self, encoder: BinaryEncoder, val: time) -> None:
-        encoder.write_time_micros_long(val)
+        encoder.write_time_micros(val)
 
 
 class TimestampWriter(Writer):
     def write(self, encoder: BinaryEncoder, val: datetime) -> None:
-        encoder.write_timestamp_micros_long(val)
+        if val.tzinfo is not None:
+            raise ValueError(f"Timestamp should not have a timezone, but has: {val.tzinfo}")
+
+        encoder.write_int(datetime_to_micros(val))
 
 
 class TimestamptzWriter(Writer):
     def write(self, encoder: BinaryEncoder, val: datetime) -> None:
-        encoder.write_timestamp_micros_long(val)
+        encoder.write_int(datetime_to_micros(val))
 
 
 class StringWriter(Writer):
@@ -103,19 +107,18 @@ def write(self, encoder: BinaryEncoder, val: Any) -> None:
 
 class UUIDWriter(Writer):
     def write(self, encoder: BinaryEncoder, val: UUID) -> None:
-        uuid_bytes = val.bytes
-
-        if len(uuid_bytes) != 16:
-            raise ValueError(f"Expected UUID to be 16 bytes, got: {len(uuid_bytes)}")
-
-        encoder.write_bytes_fixed(uuid_bytes)
+        if len(val.bytes) != 16:
+            raise ValueError(f"Expected UUID to be 16 bytes, got: {len(val.bytes)}")
+        encoder.write(val.bytes)
 
 
 @dataclass(frozen=True)
 class FixedWriter(Writer):
     _len: int = dataclassfield()
 
     def write(self, encoder: BinaryEncoder, val: bytes) -> None:
+        if len(val) != self._len:
+            raise ValueError(f"Expected {self._len} bytes, got {len(val)}")
         encoder.write(val)
 
     def __len__(self) -> int:

python/pyiceberg/expressions/literals.py

@@ -50,7 +50,7 @@
 from pyiceberg.utils.datetime import (
     date_str_to_days,
     micros_to_days,
-    time_to_micros,
+    time_str_to_micros,
     timestamp_to_micros,
     timestamptz_to_micros,
 )
@@ -558,7 +558,7 @@ def _(self, type_var: DateType) -> Literal[int]:
     @to.register(TimeType)
     def _(self, type_var: TimeType) -> Literal[int]:
         try:
-            return TimeLiteral(time_to_micros(self.value))
+            return TimeLiteral(time_str_to_micros(self.value))
         except (TypeError, ValueError) as e:
             raise ValueError(f"Could not convert {self.value} into a {type_var}") from e
 

python/pyiceberg/utils/datetime.py

@@ -61,15 +61,14 @@ def days_to_date(days: int) -> date:
     return EPOCH_DATE + timedelta(days)
 
 
-def time_to_micros(time_str: str) -> int:
+def time_str_to_micros(time_str: str) -> int:
     """Converts an ISO-8601 formatted time to microseconds from midnight."""
-    t = time.fromisoformat(time_str)
-    return (((t.hour * 60 + t.minute) * 60) + t.second) * 1_000_000 + t.microsecond
+    return time_to_micros(time.fromisoformat(time_str))
 
 
-def time_object_to_micros(t: time) -> int:
-    """Converts an datetime.time object to microseconds from midnight."""
-    return int(t.hour * 60 * 60 * 1e6 + t.minute * 60 * 1e6 + t.second * 1e6 + t.microsecond)
+def time_to_micros(t: time) -> int:
+    """Converts a datetime.time object to microseconds from midnight."""
+    return int((((t.hour * 60 * 60) + (t.minute * 60) + t.second) * 1e6) + t.microsecond)
 
 
 def datetime_to_micros(dt: datetime) -> int:
@@ -97,7 +96,7 @@ def datetime_to_millis(dt: datetime) -> int:
         delta = dt - EPOCH_TIMESTAMPTZ
     else:
         delta = dt - EPOCH_TIMESTAMP
-    return (delta.days * 86400 + delta.seconds) * 1_000 + delta.microseconds // 1_000
+    return int(delta.total_seconds() * 1_000)
 
 
 def timestamptz_to_micros(timestamptz_str: str) -> int:

python/pyiceberg/utils/schema_conversion.py

@@ -64,11 +64,11 @@
 
 LOGICAL_FIELD_TYPE_MAPPING: Dict[Tuple[str, str], PrimitiveType] = {
     ("date", "int"): DateType(),
-    ("time-millis", "int"): TimeType(),
+    ("time-micros", "long"): TimeType(),
     ("timestamp-millis", "long"): TimestampType(),
     ("time-micros", "int"): TimeType(),
-    ("timestamp-micros", "long"): TimestampType(),
     ("uuid", "fixed"): UUIDType(),
+    ("uuid", "string"): UUIDType(),
 }
 
 AvroType = Union[str, Any]
@@ -369,6 +369,11 @@ def _convert_logical_type(self, avro_logical_type: Dict[str, Any]) -> IcebergTyp
             return self._convert_logical_decimal_type(avro_logical_type)
         elif logical_type == "map":
             return self._convert_logical_map_type(avro_logical_type)
+        elif logical_type == "timestamp-micros":
+            if avro_logical_type.get("adjust-to-utc", False) is True:
+                return TimestamptzType()
+            else:
+                return TimestampType()
         elif (logical_type, physical_type) in LOGICAL_FIELD_TYPE_MAPPING:
             return LOGICAL_FIELD_TYPE_MAPPING[(logical_type, physical_type)]
         else:
@@ -542,6 +547,8 @@ def map(self, map_type: MapType, key_result: AvroType, value_result: AvroType) -
             return {
                 "type": "map",
                 "values": value_result,
+                "key-id": self.last_map_key_field_id,
+                "value-id": self.last_map_value_field_id,
             }
         else:
             # Creates a logical map that's a list of schema's
@@ -592,13 +599,13 @@ def visit_timestamp(self, timestamp_type: TimestampType) -> AvroType:
 
     def visit_timestamptz(self, timestamptz_type: TimestamptzType) -> AvroType:
         # Iceberg only supports micro's
-        return {"type": "long", "logicalType": "timestamp-micros"}
+        return {"type": "long", "logicalType": "timestamp-micros", "adjust-to-utc": True}
 
     def visit_string(self, string_type: StringType) -> AvroType:
         return "string"
 
     def visit_uuid(self, uuid_type: UUIDType) -> AvroType:
-        return {"type": "string", "logicalType": "uuid"}
+        return {"type": "fixed", "size": "16", "logicalType": "uuid"}
 
     def visit_binary(self, binary_type: BinaryType) -> AvroType:
         return "bytes"

python/tests/avro/test_decoder.py

@@ -22,7 +22,6 @@
 from io import SEEK_SET
 from types import TracebackType
 from typing import Optional, Type
-from uuid import UUID
 
 import pytest
 
@@ -180,13 +179,6 @@ def test_skip_double(decoder_class: Type[BinaryDecoder]) -> None:
     assert decoder.tell() == 8
 
 
-@pytest.mark.parametrize("decoder_class", AVAILABLE_DECODERS)
-def test_read_uuid_from_fixed(decoder_class: Type[BinaryDecoder]) -> None:
-    mis = io.BytesIO(b"\x12\x34\x56\x78" * 4)
-    decoder = decoder_class(mis)
-    assert decoder.read_uuid_from_fixed() == UUID("{12345678-1234-5678-1234-567812345678}")
-
-
 @pytest.mark.parametrize("decoder_class", AVAILABLE_DECODERS)
 def test_read_time_millis(decoder_class: Type[BinaryDecoder]) -> None:
     mis = io.BytesIO(b"\xBC\x7D")