Speed up snapshot_by_id with an index, ~4.5x faster inspect.partitions()

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."""

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"""