From 8598f6b444da675e8d1717504a131dc055c56287 Mon Sep 17 00:00:00 2001 From: ump45nose Date: Thu, 22 Oct 2026 19:01:50 +0800 Subject: [PATCH] Support reading Avro timestamp-nanos logical type --- pyiceberg/utils/schema_conversion.py | 11 ++++++++-- tests/avro/test_file.py | 27 +++++++++++++++++++++++++ tests/utils/test_schema_conversion.py | 29 +++++++++++++++++++++++++++ 3 files changed, 65 insertions(+), 2 deletions(-) diff --git a/pyiceberg/utils/schema_conversion.py b/pyiceberg/utils/schema_conversion.py index 7e6343141b..f19ff60261 100644 --- a/pyiceberg/utils/schema_conversion.py +++ b/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: diff --git a/tests/avro/test_file.py b/tests/avro/test_file.py index 175089c30b..6cc190c33b 100644 --- a/tests/avro/test_file.py +++ b/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, diff --git a/tests/utils/test_schema_conversion.py b/tests/utils/test_schema_conversion.py index b110f92bd2..3a84267680 100644 --- a/tests/utils/test_schema_conversion.py +++ b/tests/utils/test_schema_conversion.py @@ -26,6 +26,7 @@ DateType, DecimalType, FixedType, + IcebergType, IntegerType, ListType, LongType, @@ -33,7 +34,9 @@ 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"}