add TimestampNanoReader and Writer

Author: sungwy

pyiceberg/avro/reader.py

@@ -175,6 +175,14 @@ class TimestampReader(IntegerReader):
     """
 
 
+class TimestampNanoReader(IntegerReader):
+    """Reads a nanosecond granularity timestamp from the stream.
+
+    Long is decoded as python integer which represents
+    the number of nanoseconds from the unix epoch, 1 January 1970.
+    """
+
+
 class TimestamptzReader(IntegerReader):
     """Reads a microsecond granularity timestamptz from the stream.
 
@@ -185,6 +193,16 @@ class TimestamptzReader(IntegerReader):
     """
 
 
+class TimestamptzNanoReader(IntegerReader):
+    """Reads a microsecond granularity timestamptz from the stream.
+
+    Long is decoded as python integer which represents
+    the number of nanoseconds from the unix epoch, 1 January 1970.
+
+    Adjusted to UTC.
+    """
+
+
 class StringReader(Reader):
     def read(self, decoder: BinaryDecoder) -> str:
         return decoder.read_utf8()

pyiceberg/avro/resolver.py

@@ -44,7 +44,9 @@
     StringReader,
     StructReader,
     TimeReader,
+    TimestampNanoReader,
     TimestampReader,
+    TimestamptzNanoReader,
     TimestamptzReader,
     UUIDReader,
 )
@@ -63,6 +65,8 @@
     OptionWriter,
     StringWriter,
     StructWriter,
+    TimestampNanoWriter,
+    TimestamptzNanoWriter,
     TimestamptzWriter,
     TimestampWriter,
     TimeWriter,
@@ -97,7 +101,9 @@
     PrimitiveType,
     StringType,
     StructType,
+    TimestampNanoType,
     TimestampType,
+    TimestamptzNanoType,
     TimestamptzType,
     TimeType,
     UUIDType,
@@ -181,9 +187,15 @@ def visit_time(self, time_type: TimeType) -> Writer:
     def visit_timestamp(self, timestamp_type: TimestampType) -> Writer:
         return TimestampWriter()
 
+    def visit_timestamp_ns(self, timestamp_ns_type: TimestampNanoType) -> Writer:
+        return TimestampNanoWriter()
+
     def visit_timestamptz(self, timestamptz_type: TimestamptzType) -> Writer:
         return TimestamptzWriter()
 
+    def visit_timestamptz_ns(self, timestamptz_ns_type: TimestamptzNanoType) -> Writer:
+        return TimestamptzNanoWriter()
+
     def visit_string(self, string_type: StringType) -> Writer:
         return StringWriter()
 
@@ -326,9 +338,15 @@ def visit_time(self, time_type: TimeType, partner: Optional[IcebergType]) -> Wri
     def visit_timestamp(self, timestamp_type: TimestampType, partner: Optional[IcebergType]) -> Writer:
         return TimestampWriter()
 
+    def visit_timestamp_ns(self, timestamp_ns_type: TimestampNanoType, partner: Optional[IcebergType]) -> Writer:
+        return TimestampNanoWriter()
+
     def visit_timestamptz(self, timestamptz_type: TimestamptzType, partner: Optional[IcebergType]) -> Writer:
         return TimestamptzWriter()
 
+    def visit_timestamptz_ns(self, timestamptz_ns_type: TimestamptzNanoType, partner: Optional[IcebergType]) -> Writer:
+        return TimestamptzNanoWriter()
+
     def visit_string(self, string_type: StringType, partner: Optional[IcebergType]) -> Writer:
         return StringWriter()
 
@@ -456,9 +474,15 @@ def visit_time(self, time_type: TimeType, partner: Optional[IcebergType]) -> Rea
     def visit_timestamp(self, timestamp_type: TimestampType, partner: Optional[IcebergType]) -> Reader:
         return TimestampReader()
 
+    def visit_timestamp_ns(self, timestamp_ns_type: TimestampNanoType, partner: Optional[IcebergType]) -> Reader:
+        return TimestampNanoReader()
+
     def visit_timestamptz(self, timestamptz_type: TimestamptzType, partner: Optional[IcebergType]) -> Reader:
         return TimestamptzReader()
 
+    def visit_timestamptz_ns(self, timestamptz_ns_type: TimestamptzNanoType, partner: Optional[IcebergType]) -> Reader:
+        return TimestamptzNanoReader()
+
     def visit_string(self, string_type: StringType, partner: Optional[IcebergType]) -> Reader:
         return StringReader()
 

pyiceberg/avro/writer.py

@@ -95,12 +95,24 @@ def write(self, encoder: BinaryEncoder, val: int) -> None:
         encoder.write_int(val)
 
 
+@dataclass(frozen=True)
+class TimestampNanoWriter(Writer):
+    def write(self, encoder: BinaryEncoder, val: int) -> None:
+        encoder.write_int(val)
+
+
 @dataclass(frozen=True)
 class TimestamptzWriter(Writer):
     def write(self, encoder: BinaryEncoder, val: int) -> None:
         encoder.write_int(val)
 
 
+@dataclass(frozen=True)
+class TimestamptzNanoWriter(Writer):
+    def write(self, encoder: BinaryEncoder, val: int) -> None:
+        encoder.write_int(val)
+
+
 @dataclass(frozen=True)
 class StringWriter(Writer):
     def write(self, encoder: BinaryEncoder, val: Any) -> None:

pyiceberg/conversions.py

@@ -53,13 +53,15 @@
     LongType,
     PrimitiveType,
     StringType,
+    TimestampNanoType,
     TimestampType,
+    TimestamptzNanoType,
     TimestamptzType,
     TimeType,
     UUIDType,
     strtobool,
 )
-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
 from pyiceberg.utils.decimal import decimal_to_bytes, unscaled_to_decimal
 
 _BOOL_STRUCT = Struct("<?")
@@ -108,7 +110,9 @@ def _(primitive_type: BooleanType, value_str: str) -> Union[int, float, str, uui
 @partition_to_py.register(DateType)
 @partition_to_py.register(TimeType)
 @partition_to_py.register(TimestampType)
+@partition_to_py.register(TimestampNanoType)
 @partition_to_py.register(TimestamptzType)
+@partition_to_py.register(TimestamptzNanoType)
 @handle_none
 def _(primitive_type: PrimitiveType, value_str: str) -> int:
     """Convert a string to an integer value.
@@ -188,12 +192,20 @@ def _(_: PrimitiveType, value: int) -> bytes:
 
 @to_bytes.register(TimestampType)
 @to_bytes.register(TimestamptzType)
-def _(_: TimestampType, value: Union[datetime, int]) -> bytes:
+def _(_: PrimitiveType, value: Union[datetime, int]) -> bytes:
     if isinstance(value, datetime):
         value = datetime_to_micros(value)
     return _LONG_STRUCT.pack(value)
 
 
+@to_bytes.register(TimestampNanoType)
+@to_bytes.register(TimestamptzNanoType)
+def _(_: PrimitiveType, value: Union[datetime, int]) -> bytes:
+    if isinstance(value, datetime):
+        value = datetime_to_nanos(value)
+    return _LONG_STRUCT.pack(value)
+
+
 @to_bytes.register(DateType)
 def _(_: DateType, value: Union[date, int]) -> bytes:
     if isinstance(value, date):
@@ -294,6 +306,8 @@ def _(_: PrimitiveType, b: bytes) -> int:
 @from_bytes.register(TimeType)
 @from_bytes.register(TimestampType)
 @from_bytes.register(TimestamptzType)
+@from_bytes.register(TimestampNanoType)
+@from_bytes.register(TimestamptzNanoType)
 def _(_: PrimitiveType, b: bytes) -> int:
     return _LONG_STRUCT.unpack(b)[0]
 

pyiceberg/io/pyarrow.py

@@ -160,7 +160,9 @@
     PrimitiveType,
     StringType,
     StructType,
+    TimestampNanoType,
     TimestampType,
+    TimestamptzNanoType,
     TimestamptzType,
     TimeType,
     UUIDType,
@@ -648,9 +650,15 @@ def visit_time(self, _: TimeType) -> pa.DataType:
     def visit_timestamp(self, _: TimestampType) -> pa.DataType:
         return pa.timestamp(unit="us")
 
+    def visit_timestamp_ns(self, _: TimestampNanoType) -> pa.DataType:
+        return pa.timestamp(unit="ns")
+
     def visit_timestamptz(self, _: TimestamptzType) -> pa.DataType:
         return pa.timestamp(unit="us", tz="UTC")
 
+    def visit_timestamptz_ns(self, _: TimestamptzNanoType) -> pa.DataType:
+        return pa.timestamp(unit="ns", tz="UTC")
+
     def visit_string(self, _: StringType) -> pa.DataType:
         return pa.large_string()
 
@@ -1851,9 +1859,15 @@ def visit_time(self, time_type: TimeType) -> str:
     def visit_timestamp(self, timestamp_type: TimestampType) -> str:
         return "INT64"
 
+    def visit_timestamp_ns(self, timestamp_type: TimestampNanoType) -> str:
+        return "INT64"
+
     def visit_timestamptz(self, timestamptz_type: TimestamptzType) -> str:
         return "INT64"
 
+    def visit_timestamptz_ns(self, timestamptz_ns_type: TimestamptzNanoType) -> str:
+        return "INT64"
+
     def visit_string(self, string_type: StringType) -> str:
         return "BYTE_ARRAY"
 

pyiceberg/schema.py

@@ -57,7 +57,9 @@
     PrimitiveType,
     StringType,
     StructType,
+    TimestampNanoType,
     TimestampType,
+    TimestamptzNanoType,
     TimestamptzType,
     TimeType,
     UUIDType,
@@ -370,9 +372,12 @@ def check_format_version_compatibility(self, format_version: int) -> None:
         Raises:
           ValueError: If the schema is not compatible for the format version.
         """
-        for field in self._lazy_id_to_field().values():
+        for field in self._lazy_id_to_field.values():
             if format_version < field.field_type.minimum_format_version():
-                raise ValueError(f"{field.field_type} is only supported in {field.field_type.minimum_format_version()} or higher. Current format version is: {format_version}")
+                raise ValueError(
+                    f"{field.field_type} is only supported in {field.field_type.minimum_format_version()} or higher. Current format version is: {format_version}"
+                )
+
 
 class SchemaVisitor(Generic[T], ABC):
     def before_field(self, field: NestedField) -> None:
@@ -533,8 +538,12 @@ def primitive(self, primitive: PrimitiveType, primitive_partner: Optional[P]) ->
             return self.visit_time(primitive, primitive_partner)
         elif isinstance(primitive, TimestampType):
             return self.visit_timestamp(primitive, primitive_partner)
+        elif isinstance(primitive, TimestampNanoType):
+            return self.visit_timestamp_ns(primitive, primitive_partner)
         elif isinstance(primitive, TimestamptzType):
             return self.visit_timestamptz(primitive, primitive_partner)
+        elif isinstance(primitive, TimestamptzNanoType):
+            return self.visit_timestamptz_ns(primitive, primitive_partner)
         elif isinstance(primitive, StringType):
             return self.visit_string(primitive, primitive_partner)
         elif isinstance(primitive, UUIDType):
@@ -582,10 +591,18 @@ def visit_time(self, time_type: TimeType, partner: Optional[P]) -> T:
     def visit_timestamp(self, timestamp_type: TimestampType, partner: Optional[P]) -> T:
         """Visit a TimestampType."""
 
+    @abstractmethod
+    def visit_timestamp_ns(self, timestamp_ns_type: TimestampNanoType, partner: Optional[P]) -> T:
+        """Visit a TimestampNanoType."""
+
     @abstractmethod
     def visit_timestamptz(self, timestamptz_type: TimestamptzType, partner: Optional[P]) -> T:
         """Visit a TimestamptzType."""
 
+    @abstractmethod
+    def visit_timestamptz_ns(self, timestamptz_ns_type: TimestamptzNanoType, partner: Optional[P]) -> T:
+        """Visit a TimestamptzNanoType."""
+
     @abstractmethod
     def visit_string(self, string_type: StringType, partner: Optional[P]) -> T:
         """Visit a StringType."""
@@ -711,8 +728,12 @@ def primitive(self, primitive: PrimitiveType) -> T:
             return self.visit_time(primitive)
         elif isinstance(primitive, TimestampType):
             return self.visit_timestamp(primitive)
+        elif isinstance(primitive, TimestampNanoType):
+            return self.visit_timestamp_ns(primitive)
         elif isinstance(primitive, TimestamptzType):
             return self.visit_timestamptz(primitive)
+        elif isinstance(primitive, TimestamptzNanoType):
+            return self.visit_timestamptz_ns(primitive)
         elif isinstance(primitive, StringType):
             return self.visit_string(primitive)
         elif isinstance(primitive, UUIDType):
@@ -762,10 +783,18 @@ def visit_time(self, time_type: TimeType) -> T:
     def visit_timestamp(self, timestamp_type: TimestampType) -> T:
         """Visit a TimestampType."""
 
+    @abstractmethod
+    def visit_timestamp_ns(self, timestamp_type: TimestampNanoType) -> T:
+        """Visit a TimestampNanoType."""
+
     @abstractmethod
     def visit_timestamptz(self, timestamptz_type: TimestamptzType) -> T:
         """Visit a TimestamptzType."""
 
+    @abstractmethod
+    def visit_timestamptz_ns(self, timestamptz_ns_type: TimestamptzNanoType) -> T:
+        """Visit a TimestamptzNanoType."""
+
     @abstractmethod
     def visit_string(self, string_type: StringType) -> T:
         """Visit a StringType."""

pyiceberg/types.py

@@ -180,7 +180,7 @@ def is_struct(self) -> bool:
         return isinstance(self, StructType)
 
     def minimum_format_version(self) -> int:
-        """Minimum Iceberg format version after which this type is supported"""
+        """Minimum Iceberg format version after which this type is supported."""
         return 1
 
 

pyiceberg/utils/datetime.py

@@ -91,6 +91,35 @@ def timestamp_to_micros(timestamp_str: str) -> int:
     raise ValueError(f"Invalid timestamp without zone: {timestamp_str} (must be ISO-8601)")
 
 
+def time_str_to_nanos(time_str: str) -> int:
+    """Convert an ISO-8601 formatted time to nanoseconds from midnight."""
+    return time_to_nanos(time.fromisoformat(time_str))
+
+
+def time_to_nanos(t: time) -> int:
+    """Convert a datetime.time object to nanoseconds from midnight."""
+    return (((t.hour * 60 + t.minute) * 60) + t.second) * 1_000_000 + t.microsecond + t.nanosecond
+
+
+def datetime_to_nanos(dt: datetime) -> int:
+    """Convert a datetime to nanoseconds from 1970-01-01T00:00:00.000000000."""
+    if dt.tzinfo:
+        delta = dt - EPOCH_TIMESTAMPTZ
+    else:
+        delta = dt - EPOCH_TIMESTAMP
+    return (delta.days * 86400 + delta.seconds) * 1_000_000 + delta.microseconds + delta.nanoseconds
+
+
+def timestamp_to_nanos(timestamp_str: str) -> int:
+    """Convert an ISO-9601 formatted timestamp without zone to microseconds from 1970-01-01T00:00:00.000000000."""
+    if ISO_TIMESTAMP.fullmatch(timestamp_str):
+        return datetime_to_nanos(datetime.fromisoformat(timestamp_str))
+    if ISO_TIMESTAMPTZ.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}")
+    raise ValueError(f"Invalid timestamp without zone: {timestamp_str} (must be ISO-8601)")
+
+
 def datetime_to_millis(dt: datetime) -> int:
     """Convert a datetime to milliseconds from 1970-01-01T00:00:00.000000."""
     if dt.tzinfo:

pyiceberg/utils/schema_conversion.py

@@ -603,13 +603,17 @@ def visit_time(self, time_type: TimeType) -> AvroType:
         return {"type": "long", "logicalType": "time-micros"}
 
     def visit_timestamp(self, timestamp_type: TimestampType) -> AvroType:
-        # Iceberg only supports micro's
         return {"type": "long", "logicalType": "timestamp-micros", "adjust-to-utc": False}
 
+    def visit_timestamp_ns(self, timestamp_type: TimestampType) -> AvroType:
+        return {"type": "long", "logicalType": "timestamp-nanos", "adjust-to-utc": False}
+
     def visit_timestamptz(self, timestamptz_type: TimestamptzType) -> AvroType:
-        # Iceberg only supports micro's
         return {"type": "long", "logicalType": "timestamp-micros", "adjust-to-utc": True}
 
+    def visit_timestamptz_ns(self, timestamptz_type: TimestamptzType) -> AvroType:
+        return {"type": "long", "logicalType": "timestamp-nanos", "adjust-to-utc": True}
+
     def visit_string(self, string_type: StringType) -> AvroType:
         return "string"