changing approach to avoid errors and keeping the snapshot_by_id the same

Author: MatheusFreitas25

pyiceberg/table/inspect.py

@@ -310,13 +310,15 @@ def partitions(
         )
 
         partitions_map: dict[tuple[str, Any], Any] = {}
+        # snapshot_by_id is a linear scan, and there is one lookup per manifest entry
+        snapshots_by_id = {snapshot.snapshot_id: snapshot for snapshot in self.tbl.metadata.snapshots}
 
         for entry in itertools.chain.from_iterable(scan._plan_manifest_entries()):
             partition = entry.data_file.partition
             partition_record_dict = {
                 field.name: partition[pos] for pos, field in enumerate(self.tbl.metadata.specs()[entry.data_file.spec_id].fields)
             }
-            entry_snapshot = self.tbl.snapshot_by_id(entry.snapshot_id) if entry.snapshot_id is not None else None
+            entry_snapshot = snapshots_by_id.get(entry.snapshot_id) if entry.snapshot_id is not None else None
             self._update_partitions_map_from_manifest_entry(
                 partitions_map, entry.data_file, partition_record_dict, entry_snapshot
             )
@@ -532,9 +534,11 @@ def history(self) -> pa.Table:
 
         history = []
         metadata = self.tbl.metadata
+        # snapshot_by_id is a linear scan, and there is one lookup per snapshot log entry
+        snapshots_by_id = {snapshot.snapshot_id: snapshot for snapshot in metadata.snapshots}
 
         for snapshot_entry in metadata.snapshot_log:
-            snapshot = metadata.snapshot_by_id(snapshot_entry.snapshot_id)
+            snapshot = snapshots_by_id.get(snapshot_entry.snapshot_id)
 
             history.append(
                 {

pyiceberg/table/metadata.py

@@ -236,30 +236,9 @@ 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_position(self) -> dict[int, int]:
-        """Return an index of snapshot ID to its position in the snapshots list.
-
-        This is calculated once per snapshots list and cached.
-        """
-        cached = self.__dict__.get("_id_to_snapshot_position")
-        if cached is None or cached[0] is not self.snapshots:
-            cached = (
-                self.snapshots,
-                {snapshot.snapshot_id: position for position, snapshot in enumerate(self.snapshots)},
-            )
-            # The model is frozen, so bypass the pydantic __setattr__ to memoize.
-            object.__setattr__(self, "_id_to_snapshot_position", cached)
-        return cached[1]
-
     def snapshot_by_id(self, snapshot_id: int) -> Snapshot | None:
         """Get the snapshot by snapshot_id."""
-        snapshots = self.snapshots
-        if (position := self._lazy_id_to_snapshot_position.get(snapshot_id)) is not None and position < len(snapshots):
-            if (snapshot := snapshots[position]).snapshot_id == snapshot_id:
-                return snapshot
-        # Absent id, or the list was mutated in place and the cached positions no longer line up.
-        return next((snapshot for snapshot in snapshots if snapshot.snapshot_id == snapshot_id), None)
+        return next((snapshot for snapshot in self.snapshots if snapshot.snapshot_id == snapshot_id), None)
 
     def schema_by_id(self, schema_id: int) -> Schema | None:
         """Get the schema by schema_id."""

tests/table/test_metadata.py

@@ -37,7 +37,6 @@
     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
@@ -146,79 +145,6 @@ 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_reflects_in_place_snapshot_list_mutation(example_table_metadata_v2: dict[str, Any]) -> None:
-    """The snapshot lookup must not go stale when the snapshots list itself is mutated in place."""
-    table_metadata = TableMetadataV2(**example_table_metadata_v2)
-    snapshot_id = 3051729675574597004
-
-    # Build the index before mutating, so a stale one would still be in place
-    assert table_metadata.snapshot_by_id(snapshot_id) is not None
-
-    # Replacing an entry: the lookup must return the new instance, not the replaced one
-    altered = table_metadata.snapshots[0].model_copy(update={"summary": Summary(Operation.DELETE)})
-    table_metadata.snapshots[0] = altered
-    assert table_metadata.snapshot_by_id(snapshot_id) is altered
-
-    # Reordering: ids still resolve to the right snapshots
-    table_metadata.snapshots.reverse()
-    assert table_metadata.snapshot_by_id(snapshot_id) is altered
-    assert table_metadata.snapshot_by_id(3055729675574597004) is table_metadata.snapshots[0]
-
-    # Removing: the id is gone
-    table_metadata.snapshots.clear()
-    assert table_metadata.snapshot_by_id(snapshot_id) is 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_position" not in table_metadata.model_dump()
-    assert "_id_to_snapshot_position" not in table_metadata.model_dump_json()
-    assert "_id_to_snapshot_position" 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"""