fix: convert nano timestamps to nanoseconds in partition representation

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/test_transforms.py

@@ -17,7 +17,7 @@
 # under the License.
 # pylint: disable=eval-used,protected-access,redefined-outer-name
 from collections.abc import Callable
-from datetime import date
+from datetime import date, datetime
 from decimal import Decimal
 from typing import Annotated, Any
 from uuid import UUID
@@ -1625,6 +1625,30 @@ def test_strict_binary(bound_reference_binary: BoundReference) -> None:
     _assert_projection_strict(BoundIn(term=bound_reference_binary, literals={value, other_value}), transform, NotIn)
 
 
+@pytest.mark.parametrize(
+    "source_type, value, expected",
+    [
+        pytest.param(TimestampType(), None, None, id="timestamp_none"),
+        pytest.param(TimestampType(), 1577836800000000, 1577836800000000, id="timestamp_int_passthrough"),
+        pytest.param(TimestampType(), datetime(2020, 1, 1), 1577836800000000, id="timestamp_datetime"),
+        pytest.param(TimestamptzType(), datetime(2020, 1, 1), 1577836800000000, id="timestamptz_datetime"),
+        pytest.param(TimestampNanoType(), None, None, id="timestamp_nano_none"),
+        pytest.param(TimestampNanoType(), 1577836800000000000, 1577836800000000000, id="timestamp_nano_int_passthrough"),
+        pytest.param(TimestampNanoType(), datetime(2020, 1, 1), 1577836800000000000, id="timestamp_nano_datetime"),
+        pytest.param(TimestamptzNanoType(), datetime(2020, 1, 1), 1577836800000000000, id="timestamptz_nano_datetime"),
+    ],
+)
+def test_to_partition_representation_timestamps(source_type: PrimitiveType, value: Any, expected: Any) -> None:
+    assert _to_partition_representation(source_type, value) == expected
+
+
+def test_to_partition_representation_unrecognized_type_raises() -> None:
+    with pytest.raises(ValueError, match="Type not recognized"):
+        _to_partition_representation(TimestampType(), "not-a-datetime-or-int")
+    with pytest.raises(ValueError, match="Type not recognized"):
+        _to_partition_representation(TimestampNanoType(), "not-a-datetime-or-int")
+
+
 @pytest.mark.parametrize(
     "transform",
     [