add uuid partition type

Author: jqin61

pyiceberg/partitioning.py

@@ -16,6 +16,7 @@
 # under the License.
 from __future__ import annotations
 
+import uuid
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from datetime import date, datetime
@@ -53,7 +54,16 @@
     parse_transform,
 )
 from pyiceberg.typedef import IcebergBaseModel, Record
-from pyiceberg.types import DateType, IcebergType, NestedField, PrimitiveType, StructType, TimestampType, TimestamptzType
+from pyiceberg.types import (
+    DateType,
+    IcebergType,
+    NestedField,
+    PrimitiveType,
+    StructType,
+    TimestampType,
+    TimestamptzType,
+    UUIDType,
+)
 from pyiceberg.utils.datetime import date_to_days, datetime_to_micros
 
 INITIAL_PARTITION_SPEC_ID = 0
@@ -370,15 +380,15 @@ class PartitionKey:
     schema: Schema
 
     @cached_property
-    def partition(self) -> Record:  # partition key in iceberg type
+    def partition(self) -> Record:  # partition key transformed with iceberg internal representation as input
         iceberg_typed_key_values = {}
         for raw_partition_field_value in self.raw_partition_field_values:
             partition_fields = self.partition_spec.source_id_to_fields_map[raw_partition_field_value.field.source_id]
             if len(partition_fields) != 1:
                 raise ValueError("partition_fields must contain exactly one field.")
             partition_field = partition_fields[0]
             iceberg_type = self.schema.find_field(name_or_id=raw_partition_field_value.field.source_id).field_type
-            iceberg_typed_value = _to_iceberg_type(iceberg_type, raw_partition_field_value.value)
+            iceberg_typed_value = _to_iceberg_internal_representation(iceberg_type, raw_partition_field_value.value)
             transformed_value = partition_field.transform.transform(iceberg_type)(iceberg_typed_value)
             iceberg_typed_key_values[partition_field.name] = transformed_value
         return Record(**iceberg_typed_key_values)
@@ -388,21 +398,26 @@ def to_path(self) -> str:
 
 
 @singledispatch
-def _to_iceberg_type(type: IcebergType, value: Any) -> Any:
+def _to_iceberg_internal_representation(type: IcebergType, value: Any) -> Any:
     return TypeError(f"Unsupported partition field type: {type}")
 
 
-@_to_iceberg_type.register(TimestampType)
-@_to_iceberg_type.register(TimestamptzType)
+@_to_iceberg_internal_representation.register(TimestampType)
+@_to_iceberg_internal_representation.register(TimestamptzType)
 def _(type: IcebergType, value: Optional[datetime]) -> Optional[int]:
     return datetime_to_micros(value) if value is not None else None
 
 
-@_to_iceberg_type.register(DateType)
+@_to_iceberg_internal_representation.register(DateType)
 def _(type: IcebergType, value: Optional[date]) -> Optional[int]:
     return date_to_days(value) if value is not None else None
 
 
-@_to_iceberg_type.register(PrimitiveType)
+@_to_iceberg_internal_representation.register(UUIDType)
+def _(type: IcebergType, value: Optional[uuid.UUID]) -> Optional[str]:
+    return str(value) if value is not None else None
+
+
+@_to_iceberg_internal_representation.register(PrimitiveType)
 def _(type: IcebergType, value: Optional[Any]) -> Optional[Any]:
     return value

tests/integration/test_partitioning_key.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 # pylint:disable=redefined-outer-name
+import uuid
 from datetime import date, datetime, timedelta, timezone
 from decimal import Decimal
 from typing import Any, List
@@ -50,56 +51,9 @@
     StringType,
     TimestampType,
     TimestamptzType,
+    UUIDType,
 )
 
-# @pytest.fixture(scope="session")
-# def session_catalog() -> Catalog:
-#     return load_catalog(
-#         "local",
-#         **{
-#             "type": "rest",
-#             "uri": "http://localhost:8181",
-#             "s3.endpoint": "http://localhost:9000",
-#             "s3.access-key-id": "admin",
-#             "s3.secret-access-key": "password",
-#         },
-#     )
-
-
-# @pytest.fixture(scope="session")
-# def spark() -> SparkSession:
-#     import importlib.metadata
-#     import os
-
-#     spark_version = ".".join(importlib.metadata.version("pyspark").split(".")[:2])
-#     scala_version = "2.12"
-#     iceberg_version = "1.4.3"
-
-#     os.environ["PYSPARK_SUBMIT_ARGS"] = (
-#         f"--packages org.apache.iceberg:iceberg-spark-runtime-{spark_version}_{scala_version}:{iceberg_version},"
-#         f"org.apache.iceberg:iceberg-aws-bundle:{iceberg_version} pyspark-shell"
-#     )
-#     os.environ["AWS_REGION"] = "us-east-1"
-#     os.environ["AWS_ACCESS_KEY_ID"] = "admin"
-#     os.environ["AWS_SECRET_ACCESS_KEY"] = "password"
-
-#     spark = (
-#         SparkSession.builder.appName("PyIceberg integration test")
-#         .config("spark.sql.extensions", "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
-#         .config("spark.sql.catalog.integration", "org.apache.iceberg.spark.SparkCatalog")
-#         .config("spark.sql.catalog.integration.catalog-impl", "org.apache.iceberg.rest.RESTCatalog")
-#         .config("spark.sql.catalog.integration.uri", "http://localhost:8181")
-#         .config("spark.sql.catalog.integration.io-impl", "org.apache.iceberg.aws.s3.S3FileIO")
-#         .config("spark.sql.catalog.integration.warehouse", "s3://warehouse/wh/")
-#         .config("spark.sql.catalog.integration.s3.endpoint", "http://localhost:9000")
-#         .config("spark.sql.catalog.integration.s3.path-style-access", "true")
-#         .config("spark.sql.defaultCatalog", "integration")
-#         .getOrCreate()
-#     )
-
-#     return spark
-
-
 TABLE_SCHEMA = Schema(
     NestedField(field_id=1, name="boolean_field", field_type=BooleanType(), required=False),
     NestedField(field_id=2, name="string_field", field_type=StringType(), required=False),
@@ -112,10 +66,10 @@
     NestedField(field_id=9, name="timestamptz_field", field_type=TimestamptzType(), required=False),
     NestedField(field_id=10, name="date_field", field_type=DateType(), required=False),
     # NestedField(field_id=11, name="time", field_type=TimeType(), required=False),
-    # NestedField(field_id=12, name="uuid", field_type=UuidType(), required=False),
     NestedField(field_id=11, name="binary_field", field_type=BinaryType(), required=False),
     NestedField(field_id=12, name="fixed_field", field_type=FixedType(16), required=False),
-    NestedField(field_id=13, name="decimal", field_type=DecimalType(5, 2), required=False),
+    NestedField(field_id=13, name="decimal_field", field_type=DecimalType(5, 2), required=False),
+    NestedField(field_id=14, name="uuid_field", field_type=UUIDType(), required=False),
 )
 
 
@@ -353,6 +307,25 @@
             (CAST('2023-01-01' AS DATE), 'Associated string value for date 2023-01-01')
             """,
         ),
+        (
+            [PartitionField(source_id=14, field_id=1001, transform=IdentityTransform(), name="uuid_field")],
+            [uuid.UUID("f47ac10b-58cc-4372-a567-0e02b2c3d479")],
+            Record(uuid_field="f47ac10b-58cc-4372-a567-0e02b2c3d479"),
+            "uuid_field=f47ac10b-58cc-4372-a567-0e02b2c3d479",
+            f"""CREATE TABLE {identifier} (
+                uuid_field string,
+                string_field string
+            )
+            USING iceberg
+            PARTITIONED BY (
+                identity(uuid_field)
+            )
+            """,
+            f"""INSERT INTO {identifier}
+            VALUES
+            ('f47ac10b-58cc-4372-a567-0e02b2c3d479', 'Associated string value for UUID f47ac10b-58cc-4372-a567-0e02b2c3d479')
+            """,
+        ),
         (
             [PartitionField(source_id=11, field_id=1001, transform=IdentityTransform(), name="binary_field")],
             [b'example'],
@@ -770,6 +743,7 @@ def test_partition_key(
         partition_spec=spec,
         schema=TABLE_SCHEMA,
     )
+
     # key.partition is used to write the metadata in DataFile, ManifestFile and all above layers
     assert key.partition == expected_partition_record
     # key.to_path() generates the hive partitioning part of the to-write parquet file path