re-write ancestors_of() to return Iterable[Snapshot]

Author: chinmay-bhat

pyiceberg/table/__init__.py

@@ -112,6 +112,7 @@
     SnapshotLogEntry,
     SnapshotSummaryCollector,
     Summary,
+    ancestors_of,
     update_snapshot_summaries,
 )
 from pyiceberg.table.sorting import UNSORTED_SORT_ORDER, SortOrder
@@ -1292,7 +1293,16 @@ def snapshot_by_name(self, name: str) -> Optional[Snapshot]:
 
     def latest_snapshot_before_timestamp(self, timestamp_ms: int) -> Optional[Snapshot]:
         """Get the snapshot right before the given timestamp, or None if there is no matching snapshot."""
-        return self.metadata.latest_snapshot_before_timestamp(timestamp_ms)
+        result, prev_timestamp = None, 0
+        if self.metadata.current_snapshot_id is not None:
+            for snapshot in self.current_ancestors():
+                if snapshot and prev_timestamp < snapshot.timestamp_ms < timestamp_ms:
+                    result, prev_timestamp = snapshot, snapshot.timestamp_ms
+        return result
+
+    def current_ancestors(self) -> List[Optional[Snapshot]]:
+        """Get a list of ancestors of and including the current snapshot."""
+        return list(ancestors_of(self.current_snapshot(), self.metadata))  # type: ignore
 
     def history(self) -> List[SnapshotLogEntry]:
         """Get the snapshot history of this table."""

pyiceberg/table/metadata.py

@@ -230,28 +230,6 @@ def snapshot_by_id(self, snapshot_id: int) -> Optional[Snapshot]:
         """Get the snapshot by snapshot_id."""
         return next((snapshot for snapshot in self.snapshots if snapshot.snapshot_id == snapshot_id), None)
 
-    def latest_snapshot_before_timestamp(self, timestamp_ms: int) -> Optional[Snapshot]:
-        """Get the snapshot right before the given timestamp."""
-        result, prev_timestamp = None, 0
-        if self.current_snapshot_id is not None:
-            for snapshot_id, snapshot_timestamp in self.ancestors_of(self.current_snapshot_id):
-                snapshot = self.snapshot_by_id(snapshot_id)
-                if prev_timestamp < snapshot_timestamp < timestamp_ms:
-                    result, prev_timestamp = snapshot, snapshot_timestamp
-        return result
-
-    def ancestors_of(self, snapshot_id: int) -> List[tuple[int, int]]:
-        """Get the snapshot_id of the ancestors of the given snapshot."""
-        current_id: Optional[int] = snapshot_id
-        result = []
-        while current_id is not None:
-            snapshot = self.snapshot_by_id(current_id)
-            if not snapshot:
-                break
-            result.append((current_id, snapshot.timestamp_ms))
-            current_id = snapshot.parent_snapshot_id
-        return result
-
     def schema_by_id(self, schema_id: int) -> Optional[Schema]:
         """Get the schema by schema_id."""
         return next((schema for schema in self.schemas if schema.schema_id == schema_id), None)

pyiceberg/table/snapshots.py

@@ -14,17 +14,22 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from __future__ import annotations
+
 import time
 from collections import defaultdict
 from enum import Enum
-from typing import Any, DefaultDict, Dict, List, Mapping, Optional
+from typing import TYPE_CHECKING, Any, DefaultDict, Dict, Iterable, List, Mapping, Optional
 
 from pydantic import Field, PrivateAttr, model_serializer
 
 from pyiceberg.io import FileIO
 from pyiceberg.manifest import DataFile, DataFileContent, ManifestFile, read_manifest_list
 from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionSpec
 from pyiceberg.schema import Schema
+
+if TYPE_CHECKING:
+    from pyiceberg.table.metadata import TableMetadata
 from pyiceberg.typedef import IcebergBaseModel
 
 ADDED_DATA_FILES = 'added-data-files'
@@ -412,3 +417,12 @@ def _update_totals(total_property: str, added_property: str, removed_property: s
 def set_when_positive(properties: Dict[str, str], num: int, property_name: str) -> None:
     if num > 0:
         properties[property_name] = str(num)
+
+
+def ancestors_of(current_snapshot: Snapshot, table_metadata: TableMetadata) -> Iterable[Snapshot]:
+    """Get the ancestors of and including the given snapshot."""
+    if current_snapshot:
+        yield current_snapshot
+    if current_snapshot.parent_snapshot_id:
+        if parent := table_metadata.snapshot_by_id(current_snapshot.parent_snapshot_id):
+            yield from ancestors_of(parent, table_metadata)