fix(iceberg): Introduce new data sequence for RewriteFilesAction (#51)

Li0k · Li0k · commit ccb48c38d0dd · 2025-11-12T20:14:59.000+08:00
* feat(iceberg): rewrite_files support use_starting_sequence_number

* chore(test): add test_sequence_number_in_manifest_entry
diff --git a/crates/iceberg/src/scan/context.rs b/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;
@@ -348,57 +349,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),
-    )
-}
diff --git a/crates/iceberg/src/spec/snapshot.rs b/crates/iceberg/src/spec/snapshot.rs
@@ -193,6 +193,7 @@ impl Snapshot {
     }
 
     /// Get parent snapshot.
+    #[cfg(test)]
     pub(crate) fn parent_snapshot(&self, table_metadata: &TableMetadata) -> Option<SnapshotRef> {
         match self.parent_snapshot_id {
             Some(id) => table_metadata.snapshot_by_id(id).cloned(),
@@ -510,33 +511,6 @@ impl SnapshotRetention {
     }
 }
 
-/// An iterator over the ancestors of a snapshot.
-pub struct AncestorIterator<'a> {
-    current: Option<SnapshotRef>,
-    table_metadata: &'a TableMetadata,
-}
-
-impl Iterator for AncestorIterator<'_> {
-    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;
diff --git a/crates/iceberg/src/transaction/remove_snapshots.rs b/crates/iceberg/src/transaction/remove_snapshots.rs
@@ -24,9 +24,10 @@ use async_trait::async_trait;
 use itertools::Itertools;
 
 use crate::error::Result;
-use crate::spec::{MAIN_BRANCH, SnapshotReference, SnapshotRetention, TableMetadata, ancestors_of};
+use crate::spec::{MAIN_BRANCH, SnapshotReference, SnapshotRetention, TableMetadataRef};
 use crate::table::Table;
 use crate::transaction::{ActionCommit, TransactionAction};
+use crate::utils::ancestors_of;
 use crate::{Error, ErrorKind, TableRequirement, TableUpdate};
 
 /// Default value for max snapshot age in milliseconds.
@@ -103,7 +104,7 @@ impl RemoveSnapshotAction {
     fn compute_retained_refs(
         &self,
         snapshot_refs: &HashMap<String, SnapshotReference>,
-        table_meta: &TableMetadata,
+        table_meta: &TableMetadataRef,
     ) -> HashMap<String, SnapshotReference> {
         let mut retained_refs = HashMap::new();
 
@@ -144,7 +145,7 @@ impl RemoveSnapshotAction {
     fn compute_all_branch_snapshots_to_retain(
         &self,
         refs: impl Iterator<Item = SnapshotReference>,
-        table_meta: &TableMetadata,
+        table_meta: &TableMetadataRef,
     ) -> HashSet<i64> {
         let mut branch_snapshots_to_retain = HashSet::new();
         for snapshot_ref in refs {
@@ -192,11 +193,11 @@ impl RemoveSnapshotAction {
         snapshot_id: i64,
         expire_snapshots_older_than: i64,
         min_snapshots_to_keep: usize,
-        table_meta: &TableMetadata,
+        table_meta: &TableMetadataRef,
     ) -> HashSet<i64> {
         let mut ids_to_retain = HashSet::new();
         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
@@ -214,15 +215,15 @@ impl RemoveSnapshotAction {
     fn unreferenced_snapshots_to_retain(
         &self,
         refs: impl Iterator<Item = SnapshotReference>,
-        table_meta: &TableMetadata,
+        table_meta: &TableMetadataRef,
     ) -> HashSet<i64> {
         let mut ids_to_retain = HashSet::new();
         let mut referenced_snapshots = HashSet::new();
 
         for snapshot_ref in refs {
             if snapshot_ref.is_branch() {
                 if let Some(snapshot) = table_meta.snapshot_by_id(snapshot_ref.snapshot_id) {
-                    let ancestors = ancestors_of(snapshot.clone(), table_meta);
+                    let ancestors = ancestors_of(table_meta, snapshot.snapshot_id());
                     for ancestor in ancestors {
                         referenced_snapshots.insert(ancestor.snapshot_id());
                     }
@@ -251,7 +252,7 @@ impl TransactionAction for RemoveSnapshotAction {
             return Ok(ActionCommit::new(vec![], vec![]));
         }
 
-        let table_meta = table.metadata().clone();
+        let table_meta = table.metadata_ref();
 
         let mut ids_to_retain = HashSet::new();
         let retained_refs = self.compute_retained_refs(&table_meta.refs, &table_meta);
diff --git a/crates/iceberg/src/transaction/rewrite_files.rs b/crates/iceberg/src/transaction/rewrite_files.rs
@@ -53,6 +53,8 @@ pub struct RewriteFilesAction {
     removed_data_files: Vec<DataFile>,
     removed_delete_files: Vec<DataFile>,
     snapshot_id: Option<i64>,
+
+    new_data_file_sequence_number: Option<i64>,
 }
 
 struct RewriteFilesOperation;
@@ -72,6 +74,7 @@ impl RewriteFilesAction {
             removed_data_files: Vec::new(),
             removed_delete_files: Vec::new(),
             snapshot_id: None,
+            new_data_file_sequence_number: None,
         }
     }
 
@@ -121,6 +124,7 @@ impl RewriteFilesAction {
         self.min_count_to_merge = min_count_to_merge;
         self.merge_enabled = merge_enabled;
         self.snapshot_properties = properties;
+
         self
     }
 
@@ -141,6 +145,12 @@ impl RewriteFilesAction {
         self.snapshot_id = Some(snapshot_id);
         self
     }
+
+    pub fn set_new_data_file_sequence_number(mut self, seq: i64) -> Result<Self> {
+        self.new_data_file_sequence_number = Some(seq);
+
+        Ok(self)
+    }
 }
 
 impl SnapshotProduceOperation for RewriteFilesOperation {
@@ -277,7 +287,7 @@ impl SnapshotProduceOperation for RewriteFilesOperation {
 #[async_trait::async_trait]
 impl TransactionAction for RewriteFilesAction {
     async fn commit(self: Arc<Self>, table: &Table) -> Result<ActionCommit> {
-        let snapshot_producer = SnapshotProducer::new(
+        let mut snapshot_producer = SnapshotProducer::new(
             table,
             self.commit_uuid.unwrap_or_else(Uuid::now_v7),
             self.key_metadata.clone(),
@@ -300,6 +310,10 @@ impl TransactionAction for RewriteFilesAction {
                 .await?;
         }
 
+        if let Some(seq) = self.new_data_file_sequence_number {
+            snapshot_producer.set_new_data_file_sequence_number(seq);
+        }
+
         if self.merge_enabled {
             let process =
                 MergeManifestProcess::new(self.target_size_bytes, self.min_count_to_merge);
diff --git a/crates/iceberg/src/transaction/snapshot.rs b/crates/iceberg/src/transaction/snapshot.rs
@@ -91,6 +91,8 @@ pub(crate) struct SnapshotProducer<'a> {
     // It starts from 0 and increments for each new manifest file.
     // Note: This counter is limited to the range of (0..u64::MAX).
     manifest_counter: RangeFrom<u64>,
+
+    new_data_file_sequence_number: Option<i64>,
 }
 
 impl<'a> SnapshotProducer<'a> {
@@ -127,6 +129,7 @@ impl<'a> SnapshotProducer<'a> {
             removed_data_file_paths,
             removed_delete_file_paths,
             manifest_counter: (0..),
+            new_data_file_sequence_number: None,
         }
     }
 
@@ -287,7 +290,11 @@ impl<'a> SnapshotProducer<'a> {
     }
 
     // Write manifest file for added data files and return the ManifestFile for ManifestList.
-    async fn write_added_manifest(&mut self, added_files: Vec<DataFile>) -> Result<ManifestFile> {
+    async fn write_added_manifest(
+        &mut self,
+        added_files: Vec<DataFile>,
+        data_seq: Option<i64>,
+    ) -> Result<ManifestFile> {
         if added_files.is_empty() {
             return Err(Error::new(
                 ErrorKind::PreconditionFailed,
@@ -325,7 +332,9 @@ impl<'a> SnapshotProducer<'a> {
         let manifest_entries = added_files.into_iter().map(|data_file| {
             let builder = ManifestEntry::builder()
                 .status(crate::spec::ManifestStatus::Added)
-                .data_file(data_file);
+                .data_file(data_file)
+                .sequence_number_opt(data_seq);
+
             if format_version == FormatVersion::V1 {
                 builder.snapshot_id(snapshot_id).build()
             } else {
@@ -423,13 +432,17 @@ impl<'a> SnapshotProducer<'a> {
         // Process added entries.
         if !self.added_data_files.is_empty() {
             let added_data_files = std::mem::take(&mut self.added_data_files);
-            let added_manifest = self.write_added_manifest(added_data_files).await?;
+            let added_manifest = self
+                .write_added_manifest(added_data_files, self.new_data_file_sequence_number)
+                .await?;
             manifest_files.push(added_manifest);
         }
 
         if !self.added_delete_files.is_empty() {
             let added_delete_files = std::mem::take(&mut self.added_delete_files);
-            let added_manifest = self.write_added_manifest(added_delete_files).await?;
+            let added_manifest = self
+                .write_added_manifest(added_delete_files, self.new_data_file_sequence_number)
+                .await?;
             manifest_files.push(added_manifest);
         }
 
@@ -601,6 +614,10 @@ impl<'a> SnapshotProducer<'a> {
 
         Ok(ActionCommit::new(updates, requirements))
     }
+
+    pub fn set_new_data_file_sequence_number(&mut self, new_data_file_sequence_number: i64) {
+        self.new_data_file_sequence_number = Some(new_data_file_sequence_number);
+    }
 }
 
 pub(crate) struct MergeManifestProcess {
diff --git a/crates/iceberg/src/utils.rs b/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),
+    )
+}
diff --git a/crates/integration_tests/testdata/docker-compose.yaml b/crates/integration_tests/testdata/docker-compose.yaml
diff --git a/crates/integration_tests/tests/shared_tests/rewrite_files_test.rs b/crates/integration_tests/tests/shared_tests/rewrite_files_test.rs
