Fix append to a table identity-partitioned on a uuid column

pyiceberg/io/pyarrow.py

@@ -3240,7 +3240,11 @@ def _determine_partitions(spec: PartitionSpec, schema: Schema, arrow_table: pa.T
         if full_field_name is None:
             raise ValueError(f"Could not find column name for field ID: {partition.source_id}")
         field_array = _get_field_from_arrow_table(arrow_table, full_field_name)
-        arrow_table = arrow_table.append_column(name, partition.transform.pyarrow_transform(source_field.field_type)(field_array))
+        partition_array = partition.transform.pyarrow_transform(source_field.field_type)(field_array)
+        # PyArrow has no group by or equality kernels for extension types such as arrow.uuid, so use their storage type
+        if isinstance(partition_array.type, pa.BaseExtensionType):
+            partition_array = partition_array.cast(partition_array.type.storage_type)
+        arrow_table = arrow_table.append_column(name, partition_array)
 
     unique_partition_fields = arrow_table.select(partition_fields).group_by(partition_fields).aggregate([])
 

tests/io/test_pyarrow.py

@@ -137,6 +137,7 @@
     TimestampType,
     TimestamptzType,
     TimeType,
+    UUIDType,
 )
 from tests.catalog.test_base import InMemoryCatalog
 from tests.conftest import UNIFIED_AWS_SESSION_PROPERTIES
@@ -3471,6 +3472,49 @@ def test_determine_partitions_with_view_types(transform: Transform[Any, Any]) ->
     assert sorted(written.to_pylist(), key=str) == sorted(rows, key=str)
 
 
+def test_determine_partitions_with_uuid_column() -> None:
+    schema = Schema(
+        NestedField(id=1, name="id", field_type=LongType(), required=False),
+        NestedField(id=2, name="uuid", field_type=UUIDType(), required=False),
+    )
+    spec = PartitionSpec(PartitionField(source_id=2, field_id=1000, transform=IdentityTransform(), name="uuid"))
+    arrow_table = pa.table(
+        {
+            "id": pa.array([1, 2, 3, 4], type=pa.int64()),
+            "uuid": pa.array([uuid.UUID(int=1).bytes, uuid.UUID(int=2).bytes, uuid.UUID(int=1).bytes, None], type=pa.uuid()),
+        }
+    )
+
+    partitions = list(_determine_partitions(spec, schema, arrow_table))
+
+    # The partition values are the 16 bytes of each UUID
+    rows_by_partition = {p.partition_key.partition[0]: sorted(p.arrow_table_partition["id"].to_pylist()) for p in partitions}
+    assert rows_by_partition == {uuid.UUID(int=1).bytes: [1, 3], uuid.UUID(int=2).bytes: [2], None: [4]}
+    # The written data keeps the uuid type
+    assert all(p.arrow_table_partition.schema.field("uuid").type == pa.uuid() for p in partitions)
+
+
+def test_append_uuid_to_table_partitioned_on_uuid(catalog: InMemoryCatalog) -> None:
+    schema = Schema(
+        NestedField(id=1, name="id", field_type=LongType(), required=False),
+        NestedField(id=2, name="uuid", field_type=UUIDType(), required=False),
+    )
+    spec = PartitionSpec(PartitionField(source_id=2, field_id=1000, transform=IdentityTransform(), name="uuid"))
+    catalog.create_namespace("default")
+    table = catalog.create_table("default.test_append_uuid_partitioned", schema=schema, partition_spec=spec)
+
+    arrow_table = pa.table(
+        {
+            "id": pa.array([1, 2, 3], type=pa.int64()),
+            "uuid": pa.array([uuid.UUID(int=1).bytes, uuid.UUID(int=2).bytes, None], type=pa.uuid()),
+        }
+    )
+    table.append(arrow_table)
+
+    assert sorted(table.scan().to_arrow().to_pylist(), key=lambda row: row["id"]) == arrow_table.to_pylist()
+    assert {task.file.partition[0] for task in table.scan().plan_files()} == {uuid.UUID(int=1), uuid.UUID(int=2), None}
+
+
 def test_append_view_types_to_partitioned_table(catalog: InMemoryCatalog) -> None:
     schema = Schema(
         NestedField(id=1, name="id", field_type=LongType(), required=False),