Add TimestampNs Literal

Author: rambleraptor

pyiceberg/expressions/literals.py

@@ -45,7 +45,9 @@
     IntegerType,
     LongType,
     StringType,
+    TimestampNanoType,
     TimestampType,
+    TimestamptzNanoType,
     TimestamptzType,
     TimeType,
     UUIDType,
@@ -56,11 +58,17 @@
     datetime_to_micros,
     days_to_date,
     micros_to_days,
+    micros_to_nanos,
     micros_to_timestamp,
+    nanos_to_days,
+    nanos_to_micros,
+    nanos_to_timestamp,
     time_str_to_micros,
     time_to_micros,
     timestamp_to_micros,
+    timestamp_to_nanos,
     timestamptz_to_micros,
+    timestamptz_to_nanos,
 )
 from pyiceberg.utils.decimal import decimal_to_unscaled, unscaled_to_decimal
 from pyiceberg.utils.singleton import Singleton
@@ -347,6 +355,16 @@ def _(self, _: TimestampType) -> Literal[int]:
     def _(self, _: TimestamptzType) -> Literal[int]:
         return TimestampLiteral(self.value)
 
+    @to.register(TimestampNanoType)
+    def _(self, type_var: TimestampNanoType) -> Literal[int]:
+        # The value is assumed to be in microseconds, to match the TimestampType case above
+        return TimestampLiteral(self.value).to(type_var)
+
+    @to.register(TimestamptzNanoType)
+    def _(self, type_var: TimestamptzNanoType) -> Literal[int]:
+        # The value is assumed to be in microseconds, to match the TimestamptzType case above
+        return TimestampLiteral(self.value).to(type_var)
+
     @to.register(DecimalType)
     def _(self, type_var: DecimalType) -> Literal[Decimal]:
         unscaled = Decimal(self.value)
@@ -495,6 +513,54 @@ def _(self, _: TimestamptzType) -> Literal[int]:
     def _(self, _: DateType) -> Literal[int]:
         return DateLiteral(micros_to_days(self.value))
 
+    @to.register(TimestampNanoType)
+    def _(self, _: TimestampNanoType) -> Literal[int]:
+        return TimestampNanoLiteral(micros_to_nanos(self.value))
+
+    @to.register(TimestamptzNanoType)
+    def _(self, _: TimestamptzNanoType) -> Literal[int]:
+        return TimestampNanoLiteral(micros_to_nanos(self.value))
+
+
+class TimestampNanoLiteral(Literal[int]):
+    def __init__(self, value: int) -> None:
+        super().__init__(value, int)
+
+    @model_serializer
+    def ser_model(self) -> str:
+        # Python datetime only goes down to microseconds, so the last three digits are appended
+        return f"{nanos_to_timestamp(self.root).isoformat(timespec='microseconds')}{self.root % 1000:03d}"
+
+    def increment(self) -> Literal[int]:
+        return TimestampNanoLiteral(self.value + 1)
+
+    def decrement(self) -> Literal[int]:
+        return TimestampNanoLiteral(self.value - 1)
+
+    @singledispatchmethod
+    def to(self, type_var: IcebergType) -> Literal:  # type: ignore
+        raise TypeError(f"Cannot convert TimestampNanoLiteral into {type_var}")
+
+    @to.register(TimestampNanoType)
+    def _(self, _: TimestampNanoType) -> Literal[int]:
+        return self
+
+    @to.register(TimestamptzNanoType)
+    def _(self, _: TimestamptzNanoType) -> Literal[int]:
+        return self
+
+    @to.register(TimestampType)
+    def _(self, _: TimestampType) -> Literal[int]:
+        return TimestampLiteral(nanos_to_micros(self.value))
+
+    @to.register(TimestamptzType)
+    def _(self, _: TimestamptzType) -> Literal[int]:
+        return TimestampLiteral(nanos_to_micros(self.value))
+
+    @to.register(DateType)
+    def _(self, _: DateType) -> Literal[int]:
+        return DateLiteral(nanos_to_days(self.value))
+
 
 class DecimalLiteral(Literal[Decimal]):
     def __init__(self, value: Decimal) -> None:
@@ -615,6 +681,14 @@ def _(self, _: TimestampType) -> Literal[int]:
     def _(self, _: TimestamptzType) -> Literal[int]:
         return TimestampLiteral(timestamptz_to_micros(self.value))
 
+    @to.register(TimestampNanoType)
+    def _(self, _: TimestampNanoType) -> Literal[int]:
+        return TimestampNanoLiteral(timestamp_to_nanos(self.value))
+
+    @to.register(TimestamptzNanoType)
+    def _(self, _: TimestamptzNanoType) -> Literal[int]:
+        return TimestampNanoLiteral(timestamptz_to_nanos(self.value))
+
     @to.register(UUIDType)
     def _(self, _: UUIDType) -> Literal[bytes]:
         return UUIDLiteral(UUID(self.value).bytes)

pyiceberg/transforms.py

@@ -65,6 +65,7 @@
     Literal,
     LongLiteral,
     TimestampLiteral,
+    TimestampNanoLiteral,
     literal,
 )
 from pyiceberg.typedef import IcebergRootModel, L
@@ -1049,7 +1050,7 @@ def _truncate_number(
 ) -> UnboundPredicate | None:
     boundary = pred.literal
 
-    if not isinstance(boundary, (LongLiteral, DecimalLiteral, DateLiteral, TimestampLiteral)):
+    if not isinstance(boundary, (LongLiteral, DecimalLiteral, DateLiteral, TimestampLiteral, TimestampNanoLiteral)):
         raise ValueError(f"Expected a numeric literal, got: {type(boundary)}")
 
     if isinstance(pred, BoundLessThan):
@@ -1071,7 +1072,7 @@ def _truncate_number_strict(
 ) -> UnboundPredicate | None:
     boundary = pred.literal
 
-    if not isinstance(boundary, (LongLiteral, DecimalLiteral, DateLiteral, TimestampLiteral)):
+    if not isinstance(boundary, (LongLiteral, DecimalLiteral, DateLiteral, TimestampLiteral, TimestampNanoLiteral)):
         raise ValueError(f"Expected a numeric literal, got: {type(boundary)}")
 
     if isinstance(pred, BoundLessThan):

pyiceberg/utils/datetime.py

@@ -26,6 +26,8 @@
     timedelta,
 )
 
+from pyiceberg.types import LongType
+
 EPOCH_DATE = date.fromisoformat("1970-01-01")
 EPOCH_TIMESTAMP = datetime.fromisoformat("1970-01-01T00:00:00.000000")
 ISO_TIMESTAMP = re.compile(r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(.\d{1,6})?")
@@ -35,6 +37,18 @@
 ISO_TIMESTAMPTZ_NANO = re.compile(r"(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})(.\d{1,6})?(\d{1,3})?([-+]\d{2}:\d{2})")
 
 
+def _check_nanos_range(nanos: int, source: object) -> int:
+    """Reject nanosecond timestamps that do not fit in a 64-bit integer.
+
+    Python integers are unbounded, so the overflow has to be checked explicitly. Java gets this for
+    free: both DateTimeUtil.microsToNanos and DateTimeUtil.nanosFromTimestamp are exact long
+    operations that throw an ArithmeticException instead of wrapping around.
+    """
+    if not LongType.min <= nanos <= LongType.max:
+        raise OverflowError(f"Timestamp cannot be converted to nanoseconds, out of range: {source}")
+    return nanos
+
+
 def micros_to_days(timestamp: int) -> int:
     """Convert a timestamp in microseconds to a date in days."""
     return timedelta(microseconds=timestamp).days
@@ -105,8 +119,8 @@ def time_to_nanos(t: time) -> int:
     return ((((t.hour * 60 + t.minute) * 60) + t.second) * 1_000_000 + t.microsecond) * 1_000
 
 
-def datetime_to_nanos(dt: datetime) -> int:
-    """Convert a datetime to nanoseconds from 1970-01-01T00:00:00.000000000."""
+def _unchecked_datetime_to_nanos(dt: datetime) -> int:
+    """Convert a datetime to nanoseconds without a range check, so a caller can add the sub-microsecond digits first."""
     # python datetime and time doesn't have nanoseconds support yet
     # https://github.com/python/cpython/issues/59648
     if dt.tzinfo:
@@ -116,6 +130,11 @@ def datetime_to_nanos(dt: datetime) -> int:
     return ((delta.days * 86400 + delta.seconds) * 1_000_000 + delta.microseconds) * 1_000
 
 
+def datetime_to_nanos(dt: datetime) -> int:
+    """Convert a datetime to nanoseconds from 1970-01-01T00:00:00.000000000."""
+    return _check_nanos_range(_unchecked_datetime_to_nanos(dt), dt)
+
+
 def timestamp_to_nanos(timestamp_str: str) -> int:
     """Convert an ISO-9601 formatted timestamp without zone to nanoseconds from 1970-01-01T00:00:00.000000000."""
     if match := ISO_TIMESTAMP_NANO.fullmatch(timestamp_str):
@@ -126,7 +145,8 @@ def timestamp_to_nanos(timestamp_str: str) -> int:
         ns_str = (match.group(3) or "0").ljust(3, "0")
         ms_str = match.group(2) if match.group(2) else ""
         timestamp_str_without_ns_str = match.group(1) + ms_str
-        return datetime_to_nanos(datetime.fromisoformat(timestamp_str_without_ns_str)) + int(ns_str)
+        nanos = _unchecked_datetime_to_nanos(datetime.fromisoformat(timestamp_str_without_ns_str)) + int(ns_str)
+        return _check_nanos_range(nanos, timestamp_str)
     if ISO_TIMESTAMPTZ_NANO.fullmatch(timestamp_str):
         # When we can match a timestamp without a zone, we can give a more specific error
         raise ValueError(f"Zone offset provided, but not expected: {timestamp_str}")
@@ -143,7 +163,8 @@ def timestamptz_to_nanos(timestamptz_str: str) -> int:
         ns_str = (match.group(3) or "0").ljust(3, "0")
         ms_str = match.group(2) if match.group(2) else ""
         timestamptz_str_without_ns_str = match.group(1) + ms_str + match.group(4)
-        return datetime_to_nanos(datetime.fromisoformat(timestamptz_str_without_ns_str)) + int(ns_str)
+        nanos = _unchecked_datetime_to_nanos(datetime.fromisoformat(timestamptz_str_without_ns_str)) + int(ns_str)
+        return _check_nanos_range(nanos, timestamptz_str)
     if ISO_TIMESTAMP_NANO.fullmatch(timestamptz_str):
         # When we can match a timestamp without a zone, we can give a more specific error
         raise ValueError(f"Missing zone offset: {timestamptz_str} (must be ISO-8601)")
@@ -283,3 +304,8 @@ def nanos_to_hours(nanos: int) -> int:
 def nanos_to_micros(nanos: int) -> int:
     """Convert a nanoseconds timestamp to microsecond timestamp by dropping precision."""
     return nanos // 1000
+
+
+def micros_to_nanos(micros: int) -> int:
+    """Convert a microseconds timestamp to a nanosecond timestamp."""
+    return _check_nanos_range(micros * 1000, micros)

tests/expressions/test_expressions.py

@@ -77,6 +77,8 @@
     NestedField,
     StringType,
     StructType,
+    TimestampNanoType,
+    TimestamptzNanoType,
 )
 
 
@@ -110,6 +112,20 @@ def test_invert_not_nan_bind() -> None:
     assert ~NotNaN(Reference("a")).bind(schema) == IsNaN(Reference("a")).bind(schema)
 
 
+def test_bind_timestamp_nano() -> None:
+    schema = Schema(NestedField(2, "a", TimestampNanoType(), required=False), schema_id=1)
+
+    assert GreaterThan("a", "2017-08-18T14:21:01.919234567").bind(schema).literal.value == 1503066061919234567
+    # A long is read as microseconds, matching the plain timestamp case
+    assert EqualTo("a", 1503066061919234).bind(schema).literal.value == 1503066061919234000
+
+
+def test_bind_timestamptz_nano() -> None:
+    schema = Schema(NestedField(2, "a", TimestamptzNanoType(), required=False), schema_id=1)
+
+    assert GreaterThan("a", "2017-08-18T14:21:01.919234567+00:00").bind(schema).literal.value == 1503066061919234567
+
+
 def test_bind_expr_does_not_exists() -> None:
     schema = Schema(NestedField(2, "a", IntegerType()), schema_id=1)
     with pytest.raises(ValueError) as exc_info: