Support reading Avro timestamp-nanos logical type

pyiceberg/utils/schema_conversion.py

@@ -48,7 +48,9 @@
     PrimitiveType,
     StringType,
     StructType,
+    TimestampNanoType,
     TimestampType,
+    TimestamptzNanoType,
     TimestamptzType,
     TimeType,
     UnknownType,
@@ -381,6 +383,11 @@ def _convert_logical_type(self, avro_logical_type: dict[str, Any]) -> IcebergTyp
                 return TimestamptzType()
             else:
                 return TimestampType()
+        elif logical_type == "timestamp-nanos":
+            if avro_logical_type.get("adjust-to-utc", False) is True:
+                return TimestamptzNanoType()
+            else:
+                return TimestampNanoType()
         elif (logical_type, physical_type) in LOGICAL_FIELD_TYPE_MAPPING:
             return LOGICAL_FIELD_TYPE_MAPPING[(logical_type, physical_type)]
         else:
@@ -618,13 +625,13 @@ def visit_time(self, time_type: TimeType) -> AvroType:
     def visit_timestamp(self, timestamp_type: TimestampType) -> AvroType:
         return {"type": "long", "logicalType": "timestamp-micros", "adjust-to-utc": False}
 
-    def visit_timestamp_ns(self, timestamp_type: TimestampType) -> AvroType:
+    def visit_timestamp_ns(self, timestamp_ns_type: TimestampNanoType) -> AvroType:
         return {"type": "long", "logicalType": "timestamp-nanos", "adjust-to-utc": False}
 
     def visit_timestamptz(self, timestamptz_type: TimestamptzType) -> AvroType:
         return {"type": "long", "logicalType": "timestamp-micros", "adjust-to-utc": True}
 
-    def visit_timestamptz_ns(self, timestamptz_type: TimestamptzType) -> AvroType:
+    def visit_timestamptz_ns(self, timestamptz_ns_type: TimestamptzNanoType) -> AvroType:
         return {"type": "long", "logicalType": "timestamp-nanos", "adjust-to-utc": True}
 
     def visit_string(self, string_type: StringType) -> AvroType:

tests/avro/test_file.py

@@ -51,7 +51,9 @@
     LongType,
     NestedField,
     StringType,
+    TimestampNanoType,
     TimestampType,
+    TimestamptzNanoType,
     TimestamptzType,
     TimeType,
     UUIDType,
@@ -455,6 +457,31 @@ def field_uuid(self) -> UUID:
         assert record[idx] == avro_entry_read_with_fastavro[idx], f"Invalid {field} read with fastavro"
 
 
+def test_timestamp_nanos_round_trip() -> None:
+    schema = Schema(
+        NestedField(field_id=1, name="ts_ns", field_type=TimestampNanoType(), required=True),
+        NestedField(field_id=2, name="tstz_ns", field_type=TimestamptzNanoType(), required=False),
+        schema_id=1,
+    )
+    records = [
+        Record(1677629965123456789, 1677629965123456789),
+        Record(-1, None),
+    ]
+
+    with TemporaryDirectory() as tmpdir:
+        tmp_avro_file = tmpdir + "/timestamp_nanos.avro"
+        with avro.AvroOutputFile[Record](PyArrowFileIO().new_output(tmp_avro_file), schema, "timestamp_nanos") as out:
+            out.write_block(records)
+
+        with avro.AvroFile[Record](
+            PyArrowFileIO().new_input(tmp_avro_file),
+            schema,
+            {-1: Record},
+        ) as avro_reader:
+            assert avro_reader.schema == schema
+            assert list(avro_reader) == records
+
+
 def manifest_entry(index: int) -> ManifestEntry:
     return ManifestEntry.from_args(
         status=ManifestEntryStatus.ADDED,

tests/utils/test_schema_conversion.py

@@ -26,14 +26,17 @@
     DateType,
     DecimalType,
     FixedType,
+    IcebergType,
     IntegerType,
     ListType,
     LongType,
     MapType,
     NestedField,
     StringType,
     StructType,
+    TimestampNanoType,
     TimestampType,
+    TimestamptzNanoType,
     UnknownType,
     UUIDType,
 )
@@ -347,6 +350,32 @@ def test_convert_timestamp_micros_type() -> None:
     assert actual == TimestampType()
 
 
+@pytest.mark.parametrize(
+    "adjust_to_utc, expected",
+    [
+        (None, TimestampNanoType()),
+        (False, TimestampNanoType()),
+        (True, TimestamptzNanoType()),
+    ],
+)
+def test_convert_timestamp_nanos_type(adjust_to_utc: bool | None, expected: IcebergType) -> None:
+    avro_logical_type: dict[str, Any] = {"type": "long", "logicalType": "timestamp-nanos"}
+    if adjust_to_utc is not None:
+        avro_logical_type["adjust-to-utc"] = adjust_to_utc
+    actual = AvroSchemaConversion()._convert_logical_type(avro_logical_type)
+    assert actual == expected
+
+
+def test_timestamp_nanos_avro_schema_round_trip() -> None:
+    schema = Schema(
+        NestedField(field_id=1, name="ts_ns", field_type=TimestampNanoType(), required=True),
+        NestedField(field_id=2, name="tstz_ns", field_type=TimestamptzNanoType(), required=False),
+    )
+    avro_schema = AvroSchemaConversion().iceberg_to_avro(schema)
+    assert isinstance(avro_schema, dict)
+    assert AvroSchemaConversion().avro_to_iceberg(avro_schema) == schema
+
+
 def test_unknown_logical_type() -> None:
     """Test raising a ValueError when converting an unknown logical type as part of an Avro schema conversion"""
     avro_logical_type = {"type": "bytes", "logicalType": "date"}