fix(partitioning): convert nano timestamp partition values to nanoseconds

pyiceberg/partitioning.py

@@ -54,13 +54,15 @@
     NestedField,
     PrimitiveType,
     StructType,
+    TimestampNanoType,
     TimestampType,
+    TimestamptzNanoType,
     TimestamptzType,
     TimeType,
     UnknownType,
     UUIDType,
 )
-from pyiceberg.utils.datetime import date_to_days, datetime_to_micros, time_to_micros
+from pyiceberg.utils.datetime import date_to_days, datetime_to_micros, datetime_to_nanos, time_to_micros
 
 INITIAL_PARTITION_SPEC_ID = 0
 PARTITION_FIELD_ID_START: int = 1000
@@ -516,6 +518,19 @@ def _(type: IcebergType, value: int | datetime | None) -> int | None:
         raise ValueError(f"Type not recognized: {value}")
 
 
+@_to_partition_representation.register(TimestampNanoType)
+@_to_partition_representation.register(TimestamptzNanoType)
+def _(type: IcebergType, value: int | datetime | None) -> int | None:
+    if value is None:
+        return None
+    elif isinstance(value, int):
+        return value
+    elif isinstance(value, datetime):
+        return datetime_to_nanos(value)
+    else:
+        raise ValueError(f"Type not recognized: {value}")
+
+
 @_to_partition_representation.register(DateType)
 def _(type: IcebergType, value: int | date | None) -> int | None:
     if value is None:

tests/table/test_partitioning.py

@@ -22,7 +22,12 @@
 import pytest
 
 from pyiceberg.exceptions import ValidationError
-from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionField, PartitionSpec
+from pyiceberg.partitioning import (
+    UNPARTITIONED_PARTITION_SPEC,
+    PartitionField,
+    PartitionSpec,
+    _to_partition_representation,
+)
 from pyiceberg.schema import Schema
 from pyiceberg.transforms import (
     BucketTransform,
@@ -45,7 +50,9 @@
     PrimitiveType,
     StringType,
     StructType,
+    TimestampNanoType,
     TimestampType,
+    TimestamptzNanoType,
     TimestamptzType,
     TimeType,
     UnknownType,
@@ -253,6 +260,30 @@ def test_transform_consistency_with_pyarrow_transform(source_type: PrimitiveType
             assert t.transform(source_type)(value) == t.pyarrow_transform(source_type)(pa.array([value])).to_pylist()[0]
 
 
+@pytest.mark.parametrize(
+    "timestamp_type",
+    [TimestampType(), TimestamptzType()],
+)
+def test_to_partition_representation_timestamp(timestamp_type: PrimitiveType) -> None:
+    value = datetime.datetime(2020, 1, 1, 12, 30, 45, 123456)
+
+    assert _to_partition_representation(timestamp_type, value) == 1577881845123456
+    assert _to_partition_representation(timestamp_type, 1577881845123456) == 1577881845123456
+    assert _to_partition_representation(timestamp_type, None) is None
+
+
+@pytest.mark.parametrize(
+    "timestamp_nano_type",
+    [TimestampNanoType(), TimestamptzNanoType()],
+)
+def test_to_partition_representation_timestamp_nano(timestamp_nano_type: PrimitiveType) -> None:
+    value = datetime.datetime(2020, 1, 1, 12, 30, 45, 123456)
+
+    assert _to_partition_representation(timestamp_nano_type, value) == 1577881845123456000
+    assert _to_partition_representation(timestamp_nano_type, 1577881845123456789) == 1577881845123456789
+    assert _to_partition_representation(timestamp_nano_type, None) is None
+
+
 def test_deserialize_partition_field_v2() -> None:
     json_partition_spec = """{"source-id": 1, "field-id": 1000, "transform": "truncate[19]", "name": "str_truncate"}"""