From e5340126d328a0476f127c4df4dd24f3384894bf Mon Sep 17 00:00:00 2001 From: Vishnu Prakash Date: Tue, 14 Jul 2026 17:01:23 +0530 Subject: [PATCH] fix(partitioning): convert nano timestamp partition values to nanoseconds --- pyiceberg/partitioning.py | 17 +++++++++++++++- tests/table/test_partitioning.py | 33 +++++++++++++++++++++++++++++++- 2 files changed, 48 insertions(+), 2 deletions(-) diff --git a/pyiceberg/partitioning.py b/pyiceberg/partitioning.py index 3074c30ea1..bb383b69f7 100644 --- a/pyiceberg/partitioning.py +++ b/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: diff --git a/tests/table/test_partitioning.py b/tests/table/test_partitioning.py index b150fc2f67..58ada0dad4 100644 --- a/tests/table/test_partitioning.py +++ b/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"}"""