refactor(snapshot): unify ancestors (#49)

Author: Li0k

crates/iceberg/src/scan/context.rs

@@ -33,6 +33,7 @@ use crate::spec::{
     DataContentType, ManifestContentType, ManifestEntryRef, ManifestFile, ManifestList,
     ManifestStatus, Operation, SchemaRef, SnapshotRef, TableMetadataRef,
 };
+use crate::utils::ancestors_between;
 use crate::{Error, ErrorKind, Result};
 
 type ManifestEntryFilterFn = dyn Fn(&ManifestEntryRef) -> bool + Send + Sync;
@@ -350,57 +351,3 @@ impl PlanContext {
         }
     }
 }
-
-struct Ancestors {
-    next: Option<SnapshotRef>,
-    get_snapshot: Box<dyn Fn(i64) -> Option<SnapshotRef> + Send>,
-}
-
-impl Iterator for Ancestors {
-    type Item = SnapshotRef;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        let snapshot = self.next.take()?;
-        let result = snapshot.clone();
-        self.next = snapshot
-            .parent_snapshot_id()
-            .and_then(|id| (self.get_snapshot)(id));
-        Some(result)
-    }
-}
-
-/// Iterate starting from `snapshot` (inclusive) to the root snapshot.
-fn ancestors_of(
-    table_metadata: &TableMetadataRef,
-    snapshot: i64,
-) -> Box<dyn Iterator<Item = SnapshotRef> + Send> {
-    if let Some(snapshot) = table_metadata.snapshot_by_id(snapshot) {
-        let table_metadata = table_metadata.clone();
-        Box::new(Ancestors {
-            next: Some(snapshot.clone()),
-            get_snapshot: Box::new(move |id| table_metadata.snapshot_by_id(id).cloned()),
-        })
-    } else {
-        Box::new(std::iter::empty())
-    }
-}
-
-/// Iterate starting from `snapshot` (inclusive) to `oldest_snapshot_id` (exclusive).
-fn ancestors_between(
-    table_metadata: &TableMetadataRef,
-    latest_snapshot_id: i64,
-    oldest_snapshot_id: Option<i64>,
-) -> Box<dyn Iterator<Item = SnapshotRef> + Send> {
-    let Some(oldest_snapshot_id) = oldest_snapshot_id else {
-        return Box::new(ancestors_of(table_metadata, latest_snapshot_id));
-    };
-
-    if latest_snapshot_id == oldest_snapshot_id {
-        return Box::new(std::iter::empty());
-    }
-
-    Box::new(
-        ancestors_of(table_metadata, latest_snapshot_id)
-            .take_while(move |snapshot| snapshot.snapshot_id() != oldest_snapshot_id),
-    )
-}

crates/iceberg/src/spec/snapshot.rs

@@ -175,7 +175,7 @@ impl Snapshot {
     }
 
     /// Get parent snapshot.
-    pub(crate) fn parent_snapshot(&self, table_metadata: &TableMetadata) -> Option<SnapshotRef> {
+    pub fn parent_snapshot(&self, table_metadata: &TableMetadata) -> Option<SnapshotRef> {
         match self.parent_snapshot_id {
             Some(id) => table_metadata.snapshot_by_id(id).cloned(),
             None => None,
@@ -400,33 +400,6 @@ impl SnapshotRetention {
     }
 }
 
-/// An iterator over the ancestors of a snapshot.
-pub struct AncestorIterator<'a> {
-    current: Option<SnapshotRef>,
-    table_metadata: &'a TableMetadata,
-}
-
-impl<'a> Iterator for AncestorIterator<'a> {
-    type Item = SnapshotRef;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        let current = self.current.take()?;
-
-        let next = current.parent_snapshot(self.table_metadata);
-        self.current = next;
-
-        Some(current)
-    }
-}
-
-/// Returns an iterator over the ancestors of a snapshot.
-pub fn ancestors_of(snapshot: SnapshotRef, table_metadata: &TableMetadata) -> AncestorIterator<'_> {
-    AncestorIterator {
-        current: Some(snapshot),
-        table_metadata,
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use std::collections::HashMap;

crates/iceberg/src/transaction/remove_snapshots.rs

@@ -23,11 +23,12 @@ use itertools::Itertools;
 
 use crate::error::Result;
 use crate::spec::{
-    ancestors_of, SnapshotReference, SnapshotRetention, MAIN_BRANCH, MAX_REF_AGE_MS,
-    MAX_REF_AGE_MS_DEFAULT, MAX_SNAPSHOT_AGE_MS, MAX_SNAPSHOT_AGE_MS_DEFAULT,
-    MIN_SNAPSHOTS_TO_KEEP, MIN_SNAPSHOTS_TO_KEEP_DEFAULT,
+    SnapshotReference, SnapshotRetention, MAIN_BRANCH, MAX_REF_AGE_MS, MAX_REF_AGE_MS_DEFAULT,
+    MAX_SNAPSHOT_AGE_MS, MAX_SNAPSHOT_AGE_MS_DEFAULT, MIN_SNAPSHOTS_TO_KEEP,
+    MIN_SNAPSHOTS_TO_KEEP_DEFAULT,
 };
 use crate::transaction::Transaction;
+use crate::utils::ancestors_of;
 use crate::{Error, ErrorKind, TableRequirement, TableUpdate};
 
 /// RemoveSnapshotAction is a transaction action for removing snapshot.
@@ -345,9 +346,9 @@ impl<'a> RemoveSnapshotAction<'a> {
         min_snapshots_to_keep: usize,
     ) -> HashSet<i64> {
         let mut ids_to_retain = HashSet::new();
-        let table_meta = self.tx.current_table.metadata();
+        let table_meta = self.tx.current_table.metadata_ref();
         if let Some(snapshot) = table_meta.snapshot_by_id(snapshot_id) {
-            let ancestors = ancestors_of(snapshot.clone(), table_meta);
+            let ancestors = ancestors_of(&table_meta, snapshot.snapshot_id());
             for ancestor in ancestors {
                 if ids_to_retain.len() < min_snapshots_to_keep
                     || ancestor.timestamp_ms() >= expire_snapshots_older_than
@@ -377,8 +378,10 @@ impl<'a> RemoveSnapshotAction<'a> {
                     .metadata()
                     .snapshot_by_id(snapshot_ref.snapshot_id)
                 {
-                    let ancestors =
-                        ancestors_of(snapshot.clone(), self.tx.current_table.metadata());
+                    let ancestors = ancestors_of(
+                        &self.tx.current_table.metadata_ref(),
+                        snapshot.snapshot_id(),
+                    );
                     for ancestor in ancestors {
                         referenced_snapshots.insert(ancestor.snapshot_id());
                     }

crates/iceberg/src/utils.rs

@@ -17,6 +17,8 @@
 
 use std::num::NonZeroUsize;
 
+use crate::spec::{SnapshotRef, TableMetadataRef};
+
 // Use a default value of 1 as the safest option.
 // See https://doc.rust-lang.org/std/thread/fn.available_parallelism.html#limitations
 // for more details.
@@ -185,3 +187,57 @@ pub mod bin {
         }
     }
 }
+
+pub struct Ancestors {
+    next: Option<SnapshotRef>,
+    get_snapshot: Box<dyn Fn(i64) -> Option<SnapshotRef> + Send>,
+}
+
+impl Iterator for Ancestors {
+    type Item = SnapshotRef;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let snapshot = self.next.take()?;
+        let result = snapshot.clone();
+        self.next = snapshot
+            .parent_snapshot_id()
+            .and_then(|id| (self.get_snapshot)(id));
+        Some(result)
+    }
+}
+
+/// Iterate starting from `snapshot` (inclusive) to the root snapshot.
+pub fn ancestors_of(
+    table_metadata: &TableMetadataRef,
+    snapshot: i64,
+) -> Box<dyn Iterator<Item = SnapshotRef> + Send> {
+    if let Some(snapshot) = table_metadata.snapshot_by_id(snapshot) {
+        let table_metadata = table_metadata.clone();
+        Box::new(Ancestors {
+            next: Some(snapshot.clone()),
+            get_snapshot: Box::new(move |id| table_metadata.snapshot_by_id(id).cloned()),
+        })
+    } else {
+        Box::new(std::iter::empty())
+    }
+}
+
+/// Iterate starting from `snapshot` (inclusive) to `oldest_snapshot_id` (exclusive).
+pub fn ancestors_between(
+    table_metadata: &TableMetadataRef,
+    latest_snapshot_id: i64,
+    oldest_snapshot_id: Option<i64>,
+) -> Box<dyn Iterator<Item = SnapshotRef> + Send> {
+    let Some(oldest_snapshot_id) = oldest_snapshot_id else {
+        return Box::new(ancestors_of(table_metadata, latest_snapshot_id));
+    };
+
+    if latest_snapshot_id == oldest_snapshot_id {
+        return Box::new(std::iter::empty());
+    }
+
+    Box::new(
+        ancestors_of(table_metadata, latest_snapshot_id)
+            .take_while(move |snapshot| snapshot.snapshot_id() != oldest_snapshot_id),
+    )
+}