Add data_sequence_number to ReplaceDataFilesAction

Author: Shekharrajak

crates/iceberg/src/transaction/replace_data_files.rs

@@ -38,6 +38,7 @@ pub struct ReplaceDataFilesAction {
     files_to_delete: Vec<DataFile>,
     files_to_add: Vec<DataFile>,
     validate_from_snapshot_id: Option<i64>,
+    data_sequence_number: Option<i64>,
 }
 
 impl ReplaceDataFilesAction {
@@ -49,6 +50,7 @@ impl ReplaceDataFilesAction {
             files_to_delete: vec![],
             files_to_add: vec![],
             validate_from_snapshot_id: None,
+            data_sequence_number: None,
         }
     }
 
@@ -87,6 +89,12 @@ impl ReplaceDataFilesAction {
         self.validate_from_snapshot_id = Some(snapshot_id);
         self
     }
+
+    /// Set data sequence number for new files (for handling concurrent equality deletes).
+    pub fn data_sequence_number(mut self, seq_num: i64) -> Self {
+        self.data_sequence_number = Some(seq_num);
+        self
+    }
 }
 
 #[async_trait]
@@ -111,14 +119,18 @@ impl TransactionAction for ReplaceDataFilesAction {
             Self::validate_files_exist(table, snapshot_id, &self.files_to_delete).await?;
         }
 
-        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(),
             self.snapshot_properties.clone(),
             self.files_to_add.clone(),
         );
 
+        if let Some(seq_num) = self.data_sequence_number {
+            snapshot_producer = snapshot_producer.with_data_sequence_number(seq_num);
+        }
+
         snapshot_producer.validate_added_data_files()?;
 
         let replace_op = ReplaceOperation {
@@ -501,4 +513,45 @@ mod integration_tests {
         let result = tx.commit(&catalog).await;
         assert!(result.is_err());
     }
+
+    #[tokio::test]
+    async fn test_data_sequence_number() {
+        let catalog = new_memory_catalog().await;
+        let table = make_v3_minimal_table_in_catalog(&catalog).await;
+
+        let file1 = create_file("data/file1.parquet", 100);
+        let tx = Transaction::new(&table);
+        let action = tx.fast_append().add_data_files(vec![file1.clone()]);
+        let tx = action.apply(tx).unwrap();
+        let table = tx.commit(&catalog).await.unwrap();
+
+        let original_seq = table.metadata().current_snapshot().unwrap().sequence_number();
+
+        // Replace with custom sequence number
+        let compacted = create_file("data/compacted.parquet", 100);
+        let tx = Transaction::new(&table);
+        let action = tx
+            .replace_data_files()
+            .data_sequence_number(original_seq)
+            .delete_files(vec![file1])
+            .add_files(vec![compacted]);
+        let tx = action.apply(tx).unwrap();
+        let table = tx.commit(&catalog).await.unwrap();
+
+        // Verify the new manifest entry has the custom sequence number
+        let snapshot = table.metadata().current_snapshot().unwrap();
+        let manifest_list = snapshot
+            .load_manifest_list(table.file_io(), table.metadata())
+            .await
+            .unwrap();
+
+        for entry in manifest_list.entries() {
+            let manifest = entry.load_manifest(table.file_io()).await.unwrap();
+            for e in manifest.entries() {
+                if e.file_path() == "data/compacted.parquet" {
+                    assert_eq!(e.sequence_number(), Some(original_seq));
+                }
+            }
+        }
+    }
 }

crates/iceberg/src/transaction/snapshot.rs

@@ -118,6 +118,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>,
+    // Custom data sequence number for compaction operations.
+    data_sequence_number: Option<i64>,
 }
 
 impl<'a> SnapshotProducer<'a> {
@@ -136,9 +138,15 @@ impl<'a> SnapshotProducer<'a> {
             snapshot_properties,
             added_data_files,
             manifest_counter: (0..),
+            data_sequence_number: None,
         }
     }
 
+    pub(crate) fn with_data_sequence_number(mut self, seq_num: i64) -> Self {
+        self.data_sequence_number = Some(seq_num);
+        self
+    }
+
     pub(crate) fn validate_added_data_files(&self) -> Result<()> {
         for data_file in &self.added_data_files {
             if data_file.content_type() != crate::spec::DataContentType::Data {
@@ -300,18 +308,28 @@ impl<'a> SnapshotProducer<'a> {
 
         let snapshot_id = self.snapshot_id;
         let format_version = self.table.metadata().format_version();
-        let manifest_entries = added_data_files.into_iter().map(|data_file| {
-            let builder = ManifestEntry::builder()
-                .status(crate::spec::ManifestStatus::Added)
-                .data_file(data_file);
-            if format_version == FormatVersion::V1 {
-                builder.snapshot_id(snapshot_id).build()
-            } else {
-                // For format version > 1, we set the snapshot id at the inherited time to avoid rewrite the manifest file when
-                // commit failed.
-                builder.build()
-            }
-        });
+        let data_sequence_number = self.data_sequence_number;
+        let manifest_entries: Vec<_> = added_data_files
+            .into_iter()
+            .map(|data_file| {
+                match (format_version, data_sequence_number) {
+                    (FormatVersion::V1, _) => ManifestEntry::builder()
+                        .status(crate::spec::ManifestStatus::Added)
+                        .snapshot_id(snapshot_id)
+                        .data_file(data_file)
+                        .build(),
+                    (_, Some(seq_num)) => ManifestEntry::builder()
+                        .status(crate::spec::ManifestStatus::Added)
+                        .sequence_number(seq_num)
+                        .data_file(data_file)
+                        .build(),
+                    (_, None) => ManifestEntry::builder()
+                        .status(crate::spec::ManifestStatus::Added)
+                        .data_file(data_file)
+                        .build(),
+                }
+            })
+            .collect();
         let mut writer = self.new_manifest_writer(ManifestContentType::Data)?;
         for entry in manifest_entries {
             writer.add_entry(entry)?;