diff --git a/pyiceberg/table/metadata.py b/pyiceberg/table/metadata.py index 26b6e3d3ad..ac5c3ca1eb 100644 --- a/pyiceberg/table/metadata.py +++ b/pyiceberg/table/metadata.py @@ -236,9 +236,19 @@ class TableMetadataCommonFields(IcebergBaseModel): def transform_properties_dict_value_to_str(cls, properties: Properties) -> dict[str, str]: return transform_dict_value_to_str(properties) + @property + def _lazy_id_to_snapshot(self) -> dict[int, Snapshot]: + """Return an index of snapshot ID to Snapshot instance.This is calculated once per snapshots list and cached.""" + cached = self.__dict__.get("_id_to_snapshot") + if cached is None or cached[0] is not self.snapshots: + cached = (self.snapshots, {snapshot.snapshot_id: snapshot for snapshot in self.snapshots}) + # The model is frozen, so bypass the pydantic __setattr__ to memoize. + object.__setattr__(self, "_id_to_snapshot", cached) + return cached[1] + def snapshot_by_id(self, snapshot_id: int) -> Snapshot | None: """Get the snapshot by snapshot_id.""" - return next((snapshot for snapshot in self.snapshots if snapshot.snapshot_id == snapshot_id), None) + return self._lazy_id_to_snapshot.get(snapshot_id) def schema_by_id(self, schema_id: int) -> Schema | None: """Get the schema by schema_id.""" diff --git a/tests/table/test_metadata.py b/tests/table/test_metadata.py index c163c90626..3583296791 100644 --- a/tests/table/test_metadata.py +++ b/tests/table/test_metadata.py @@ -37,6 +37,7 @@ new_table_metadata, ) from pyiceberg.table.refs import SnapshotRef, SnapshotRefType +from pyiceberg.table.snapshots import Operation, Snapshot, Summary from pyiceberg.table.sorting import NullOrder, SortDirection, SortField, SortOrder from pyiceberg.transforms import IdentityTransform from pyiceberg.typedef import UTF8 @@ -145,6 +146,56 @@ def test_parsing_correct_types(example_table_metadata_v2: dict[str, Any]) -> Non assert isinstance(table_metadata.schemas[0].fields[0].field_type, LongType) +def test_snapshot_by_id(example_table_metadata_v2: dict[str, Any]) -> None: + table_metadata = TableMetadataV2(**example_table_metadata_v2) + + # Returns the same instance that is in the snapshots list, not a copy + assert table_metadata.snapshot_by_id(3051729675574597004) is table_metadata.snapshots[0] + assert table_metadata.snapshot_by_id(3055729675574597004) is table_metadata.snapshots[1] + assert table_metadata.snapshot_by_id(-1) is None + + +def test_snapshot_by_id_index_is_invalidated_on_model_copy(example_table_metadata_v2: dict[str, Any]) -> None: + """The snapshot lookup index must not survive a model_copy that replaces the snapshots.""" + table_metadata = TableMetadataV2(**example_table_metadata_v2) + + # Build the index before copying, so a stale one would be carried over + assert table_metadata.snapshot_by_id(3051729675574597004) is not None + + new_snapshot = Snapshot( + snapshot_id=1, + parent_snapshot_id=3055729675574597004, + sequence_number=35, + timestamp_ms=1602638573591, + manifest_list="s3://bucket/test/manifest-list", + summary=Summary(Operation.APPEND), + schema_id=1, + ) + with_added = table_metadata.model_copy(update={"snapshots": table_metadata.snapshots + [new_snapshot]}) + assert with_added.snapshot_by_id(1) is new_snapshot + assert with_added.snapshot_by_id(3051729675574597004) is not None + + without_first = with_added.model_copy(update={"snapshots": with_added.snapshots[1:]}) + assert without_first.snapshot_by_id(3051729675574597004) is None + assert without_first.snapshot_by_id(1) is new_snapshot + + # The original is unaffected by either copy + assert table_metadata.snapshot_by_id(1) is None + assert table_metadata.snapshot_by_id(3051729675574597004) is not None + + +def test_snapshot_by_id_index_is_not_serialized(example_table_metadata_v2: dict[str, Any]) -> None: + """The memoized index is an implementation detail and must stay out of the serialized form.""" + table_metadata = TableMetadataV2(**example_table_metadata_v2) + assert table_metadata.snapshot_by_id(3051729675574597004) is not None + + assert "_id_to_snapshot" not in table_metadata.model_dump() + assert "_id_to_snapshot" not in table_metadata.model_dump_json() + assert "_id_to_snapshot" not in TableMetadataV2.model_fields + # Two equal instances stay equal when only one of them has built the index + assert table_metadata == TableMetadataV2(**example_table_metadata_v2) + + def test_updating_metadata(example_table_metadata_v2: dict[str, Any]) -> None: """Test creating a new TableMetadata instance that's an updated version of an existing TableMetadata instance"""