Add JSON single-value serialization for timestamp_ns and timestamptz_ns

Author: takayoshi-makabe

pyiceberg/conversions.py

@@ -76,11 +76,15 @@
     time_str_to_micros,
     time_to_micros,
     timestamp_to_micros,
+    timestamp_to_nanos,
     timestamptz_to_micros,
+    timestamptz_to_nanos,
     to_human_day,
     to_human_time,
     to_human_timestamp,
+    to_human_timestamp_ns,
     to_human_timestamptz,
+    to_human_timestamptz_ns,
 )
 from pyiceberg.utils.decimal import decimal_to_bytes, unscaled_to_decimal
 
@@ -463,6 +467,22 @@ def _(_: TimestamptzType, val: int | datetime) -> str:
     return to_human_timestamptz(val)
 
 
+@to_json.register(TimestampNanoType)
+def _(_: TimestampNanoType, val: int | datetime) -> str:
+    """Python datetime (without timezone) or nanoseconds since epoch serializes into an ISO8601 timestamp."""
+    if isinstance(val, datetime):
+        val = datetime_to_nanos(val)
+    return to_human_timestamp_ns(val)
+
+
+@to_json.register(TimestamptzNanoType)
+def _(_: TimestamptzNanoType, val: int | datetime) -> str:
+    """Python datetime (with timezone) or nanoseconds since epoch serializes into an ISO8601 timestamp."""
+    if isinstance(val, datetime):
+        val = datetime_to_nanos(val)
+    return to_human_timestamptz_ns(val)
+
+
 @to_json.register(FloatType)
 @to_json.register(DoubleType)
 def _(_: FloatType | DoubleType, val: float) -> float:
@@ -613,6 +633,34 @@ def _(_: TimestamptzType, val: str | int | datetime) -> datetime:
         return val
 
 
+@from_json.register(TimestampNanoType)
+def _(_: TimestampNanoType, val: str | int | datetime) -> int:
+    """JSON ISO8601 string into nanoseconds since epoch.
+
+    Python datetime cannot hold nanoseconds, so the value stays an int.
+    """
+    if isinstance(val, str):
+        return timestamp_to_nanos(val)
+    elif isinstance(val, datetime):
+        return datetime_to_nanos(val)
+    else:
+        return val
+
+
+@from_json.register(TimestamptzNanoType)
+def _(_: TimestamptzNanoType, val: str | int | datetime) -> int:
+    """JSON ISO8601 string into nanoseconds since epoch.
+
+    Python datetime cannot hold nanoseconds, so the value stays an int.
+    """
+    if isinstance(val, str):
+        return timestamptz_to_nanos(val)
+    elif isinstance(val, datetime):
+        return datetime_to_nanos(val)
+    else:
+        return val
+
+
 @from_json.register(FloatType)
 @from_json.register(DoubleType)
 def _(_: FloatType | DoubleType, val: float) -> float:

pyiceberg/utils/datetime.py

@@ -222,6 +222,20 @@ def to_human_timestamp(timestamp_micros: int) -> str:
     return (EPOCH_TIMESTAMP + timedelta(microseconds=timestamp_micros)).isoformat()
 
 
+def to_human_timestamp_ns(timestamp_nanos: int) -> str:
+    """Convert a TimestampNanoType value to human string."""
+    seconds, nanos = divmod(timestamp_nanos, 1_000_000_000)
+    timestamp = EPOCH_TIMESTAMP + timedelta(seconds=seconds)
+    return f"{timestamp.isoformat(timespec='seconds')}.{nanos:09d}"
+
+
+def to_human_timestamptz_ns(timestamp_nanos: int) -> str:
+    """Convert a TimestamptzNanoType value to human string."""
+    seconds, nanos = divmod(timestamp_nanos, 1_000_000_000)
+    timestamp = EPOCH_TIMESTAMPTZ + timedelta(seconds=seconds)
+    return f"{timestamp.replace(tzinfo=None).isoformat(timespec='seconds')}.{nanos:09d}+00:00"
+
+
 def micros_to_hours(micros: int) -> int:
     """Convert a timestamp in microseconds to hours from 1970-01-01T00:00."""
     return micros // 3_600_000_000

tests/test_conversions.py

@@ -575,6 +575,15 @@ def test_datetime_obj_to_bytes(primitive_type: PrimitiveType, value: datetime |
         (TimeType(), time(22, 31, 8, 123456), "22:31:08.123456"),
         (TimestampType(), datetime(2017, 11, 16, 22, 31, 8, 123456), "2017-11-16T22:31:08.123456"),
         (TimestamptzType(), datetime(2017, 11, 16, 22, 31, 8, 123456, tzinfo=timezone.utc), "2017-11-16T22:31:08.123456+00:00"),
+        (TimestampNanoType(), 1510871468123456789, "2017-11-16T22:31:08.123456789"),
+        (TimestamptzNanoType(), 1510871468123456789, "2017-11-16T22:31:08.123456789+00:00"),
+        # Python datetime only carries microseconds, so the last three digits are zero
+        (TimestampNanoType(), datetime(2017, 11, 16, 22, 31, 8, 123456), "2017-11-16T22:31:08.123456000"),
+        (
+            TimestamptzNanoType(),
+            datetime(2017, 11, 16, 22, 31, 8, 123456, tzinfo=timezone.utc),
+            "2017-11-16T22:31:08.123456000+00:00",
+        ),
         (StringType(), "iceberg", "iceberg"),
         (BinaryType(), b"\x01\x02\x03\xff", "010203ff"),
         (FixedType(4), b"\x01\x02\x03\xff", "010203ff"),
@@ -600,6 +609,8 @@ def test_json_single_serialization(primitive_type: PrimitiveType, value: Any, ex
         (TimeType(), time(22, 31, 8, 123456)),
         (TimestampType(), datetime(2017, 11, 16, 22, 31, 8, 123456)),
         (TimestamptzType(), datetime(2017, 11, 16, 22, 31, 8, 123456, tzinfo=timezone.utc)),
+        (TimestampNanoType(), 1510871468123456789),
+        (TimestamptzNanoType(), 1510871468123456789),
         (StringType(), "iceberg"),
         (BinaryType(), b"\x01\x02\x03\xff"),
         (FixedType(4), b"\x01\x02\x03\xff"),

tests/test_types.py

@@ -43,7 +43,9 @@
     PrimitiveType,
     StringType,
     StructType,
+    TimestampNanoType,
     TimestampType,
+    TimestamptzNanoType,
     TimestamptzType,
     TimeType,
     UUIDType,
@@ -947,3 +949,22 @@ def test_decimal_scale_validation() -> None:
 
     with pytest.raises(ValidationError, match="Decimal scale must be between 0 and the precision"):
         DecimalType(5, 10)
+
+
+@pytest.mark.parametrize(
+    "field_type, expected_json",
+    [
+        (TimestampNanoType(), "2017-11-16T22:31:08.123456789"),
+        (TimestamptzNanoType(), "2017-11-16T22:31:08.123456789+00:00"),
+    ],
+)
+def test_nested_field_nanosecond_defaults(field_type: PrimitiveType, expected_json: str) -> None:
+    """Nanosecond timestamp defaults serialize to ISO8601 and survive a round-trip."""
+    nanos = 1510871468123456789
+    field = NestedField(1, "ts", field_type, required=False, initial_default=nanos, write_default=nanos)
+
+    serialized = field.model_dump_json()
+    assert f'"initial-default":"{expected_json}"' in serialized
+    assert f'"write-default":"{expected_json}"' in serialized
+
+    assert NestedField.model_validate_json(serialized) == field

tests/utils/test_datetime.py

@@ -30,6 +30,8 @@
     time_to_nanos,
     timestamp_to_nanos,
     timestamptz_to_nanos,
+    to_human_timestamp_ns,
+    to_human_timestamptz_ns,
 )
 
 timezones = [
@@ -169,3 +171,32 @@ def test_nanos_to_micros(nanos: int, micros: int) -> None:
 )
 def test_nanos_to_hours(nanos: int, hours: int) -> None:
     assert hours == nanos_to_hours(nanos)
+
+
+@pytest.mark.parametrize(
+    "nanos, expected",
+    [
+        (0, "1970-01-01T00:00:00.000000000"),
+        (1510871468123456789, "2017-11-16T22:31:08.123456789"),
+        # sub-second digits are zero padded to nine positions
+        (1510871468000000001, "2017-11-16T22:31:08.000000001"),
+        (-1, "1969-12-31T23:59:59.999999999"),
+    ],
+)
+def test_to_human_timestamp_ns(nanos: int, expected: str) -> None:
+    assert to_human_timestamp_ns(nanos) == expected
+    assert timestamp_to_nanos(expected) == nanos
+
+
+@pytest.mark.parametrize(
+    "nanos, expected",
+    [
+        (0, "1970-01-01T00:00:00.000000000+00:00"),
+        (1510871468123456789, "2017-11-16T22:31:08.123456789+00:00"),
+        (1510871468000000001, "2017-11-16T22:31:08.000000001+00:00"),
+        (-1, "1969-12-31T23:59:59.999999999+00:00"),
+    ],
+)
+def test_to_human_timestamptz_ns(nanos: int, expected: str) -> None:
+    assert to_human_timestamptz_ns(nanos) == expected
+    assert timestamptz_to_nanos(expected) == nanos