rewrite the writer for support both bucket and identity transform

Author: dingo4dev

pyiceberg/avro/writer.py

@@ -32,6 +32,7 @@
     List,
     Optional,
     Tuple,
+    Union,
 )
 from uuid import UUID
 
@@ -121,8 +122,23 @@ def write(self, encoder: BinaryEncoder, val: Any) -> None:
 
 @dataclass(frozen=True)
 class UUIDWriter(Writer):
-    def write(self, encoder: BinaryEncoder, val: UUID) -> None:
-        encoder.write(val.bytes)
+    def write(self, encoder: BinaryEncoder, val: Union[UUID, str, bytes]) -> None:
+        if isinstance(val, UUID):
+            encoder.write(val.bytes)
+        elif isinstance(val, bytes):
+            encoder.write(val)
+        elif isinstance(val, str):
+            print(f"UUID string: {val=}")
+            if val.startswith("b'") and val.endswith("'"):
+                # Handle string representation of bytes
+                # Convert the escaped string to actual bytes
+                byte_string = val[2:-1].encode("utf-8").decode("unicode_escape").encode("latin1")
+                encoder.write(UUID(bytes=byte_string).bytes)
+            else:
+                # Regular UUID string
+                encoder.write(UUID(val).bytes)
+        else:
+            raise TypeError(f"Expected UUID, bytes, or string, got {type(val)}")
 
 
 @dataclass(frozen=True)

pyiceberg/conversions.py

@@ -267,10 +267,13 @@ def _(_: StringType, value: str) -> bytes:
 
 
 @to_bytes.register(UUIDType)
-def _(_: UUIDType, value: Union[uuid.UUID, bytes]) -> bytes:
+def _(_: UUIDType, value: Union[uuid.UUID, bytes, str]) -> bytes:
     if isinstance(value, bytes):
-        return value
-    return value.bytes
+        return str(uuid.UUID(bytes=value)).encode(UTF8)
+    elif isinstance(value, uuid.UUID):
+        return str(value).encode(UTF8)
+    else:
+        return str(uuid.UUID(value)).encode(UTF8)
 
 
 @to_bytes.register(BinaryType)
@@ -355,11 +358,15 @@ def _(_: StringType, b: bytes) -> str:
 
 @from_bytes.register(BinaryType)
 @from_bytes.register(FixedType)
-@from_bytes.register(UUIDType)
 def _(_: PrimitiveType, b: bytes) -> bytes:
     return b
 
 
+@from_bytes.register(UUIDType)
+def _(_: UUIDType, b: bytes) -> uuid.UUID:
+    return uuid.UUID(b.decode(UTF8))
+
+
 @from_bytes.register(DecimalType)
 def _(primitive_type: DecimalType, buf: bytes) -> Decimal:
     unscaled = int.from_bytes(buf, "big", signed=True)

pyiceberg/expressions/literals.py

@@ -144,7 +144,7 @@ def literal(value: L) -> Literal[L]:
     elif isinstance(value, str):
         return StringLiteral(value)
     elif isinstance(value, UUID):
-        return UUIDLiteral(value.bytes)  # type: ignore
+        return UUIDLiteral(UUID(bytes=value.bytes))
     elif isinstance(value, bytes):
         return BinaryLiteral(value)
     elif isinstance(value, Decimal):
@@ -586,8 +586,8 @@ def _(self, _: TimestamptzType) -> Literal[int]:
         return TimestampLiteral(timestamptz_to_micros(self.value))
 
     @to.register(UUIDType)
-    def _(self, _: UUIDType) -> Literal[bytes]:
-        return UUIDLiteral(UUID(self.value).bytes)
+    def _(self, _: UUIDType) -> Literal[UUID]:
+        return UUIDLiteral(UUID(self.value))
 
     @to.register(DecimalType)
     def _(self, type_var: DecimalType) -> Literal[Decimal]:
@@ -631,30 +631,30 @@ def __repr__(self) -> str:
         return f"literal({repr(self.value)})"
 
 
-class UUIDLiteral(Literal[bytes]):
-    def __init__(self, value: bytes) -> None:
-        super().__init__(value, bytes)
+class UUIDLiteral(Literal[UUID]):
+    def __init__(self, value: UUID) -> None:
+        super().__init__(value, UUID)
 
     @singledispatchmethod
     def to(self, type_var: IcebergType) -> Literal:  # type: ignore
         raise TypeError(f"Cannot convert UUIDLiteral into {type_var}")
 
     @to.register(UUIDType)
-    def _(self, _: UUIDType) -> Literal[bytes]:
+    def _(self, _: UUIDType) -> Literal[UUID]:
         return self
 
     @to.register(FixedType)
     def _(self, type_var: FixedType) -> Literal[bytes]:
         if len(type_var) == UUID_BYTES_LENGTH:
-            return FixedLiteral(self.value)
+            return FixedLiteral(self.value.bytes)
         else:
             raise TypeError(
                 f"Cannot convert UUIDLiteral into {type_var}, different length: {len(type_var)} <> {UUID_BYTES_LENGTH}"
             )
 
     @to.register(BinaryType)
     def _(self, _: BinaryType) -> Literal[bytes]:
-        return BinaryLiteral(self.value)
+        return BinaryLiteral(self.value.bytes)
 
 
 class FixedLiteral(Literal[bytes]):
@@ -679,9 +679,9 @@ def _(self, _: BinaryType) -> Literal[bytes]:
         return BinaryLiteral(self.value)
 
     @to.register(UUIDType)
-    def _(self, type_var: UUIDType) -> Literal[bytes]:
+    def _(self, type_var: UUIDType) -> Literal[UUID]:
         if len(self.value) == UUID_BYTES_LENGTH:
-            return UUIDLiteral(self.value)
+            return UUIDLiteral(UUID(bytes=self.value))
         else:
             raise TypeError(
                 f"Could not convert {self.value!r} into a {type_var}, lengths differ {len(self.value)} <> {UUID_BYTES_LENGTH}"
@@ -710,9 +710,9 @@ def _(self, type_var: FixedType) -> Literal[bytes]:
             )
 
     @to.register(UUIDType)
-    def _(self, type_var: UUIDType) -> Literal[bytes]:
+    def _(self, type_var: UUIDType) -> Literal[UUID]:
         if len(self.value) == UUID_BYTES_LENGTH:
-            return UUIDLiteral(self.value)
+            return UUIDLiteral(UUID(bytes=self.value))
         else:
             raise TypeError(
                 f"Cannot convert BinaryLiteral into {type_var}, different length: {UUID_BYTES_LENGTH} <> {len(self.value)}"

pyiceberg/partitioning.py

@@ -467,9 +467,11 @@ def _(type: IcebergType, value: Optional[time]) -> Optional[int]:
 
 
 @_to_partition_representation.register(UUIDType)
-def _(type: IcebergType, value: Optional[Union[uuid.UUID, int]]) -> Optional[Union[str, int]]:
+def _(type: IcebergType, value: Optional[Union[uuid.UUID, int, bytes]]) -> Optional[Union[str, int]]:
     if value is None:
         return None
+    elif isinstance(value, bytes):
+        return str(uuid.UUID(bytes=value))  # IdentityTransform
     elif isinstance(value, uuid.UUID):
         return str(value)  # IdentityTransform
     elif isinstance(value, int):