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

* feat(iceberg): rewrite_files support use_starting_sequence_number

* chore(test): add test_sequence_number_in_manifest_entry

Author: Li0k

crates/iceberg/src/transaction/rewrite_files.rs

@@ -34,6 +34,9 @@ use crate::spec::{
     Operation,
 };
 
+pub const USE_STARTING_SEQUENCE_NUMBER: &str = "use-starting-sequence-number";
+pub const USE_STARTING_SEQUENCE_NUMBER_DEFAULT: bool = true;
+
 /// Transaction action for rewriting files.
 pub struct RewriteFilesAction<'a> {
     snapshot_produce_action: SnapshotProduceAction<'a>,
@@ -74,15 +77,40 @@ impl<'a> RewriteFilesAction<'a> {
             .and_then(|s| s.parse().ok())
             .unwrap_or(MANIFEST_MERGE_ENABLED_DEFAULT);
 
+        // If the compaction should use the sequence number of the snapshot at compaction start time for
+        // new data files, instead of using the sequence number of the newly produced snapshot.
+        // This avoids commit conflicts with updates that add newer equality deletes at a higher sequence number.
+        let use_starting_sequence_number = tx
+            .current_table
+            .metadata()
+            .properties()
+            .get(USE_STARTING_SEQUENCE_NUMBER)
+            .and_then(|s| s.parse().ok())
+            .unwrap_or(USE_STARTING_SEQUENCE_NUMBER_DEFAULT);
+
+        let mut snapshot_produce_action = SnapshotProduceAction::new(
+            tx,
+            snapshot_id,
+            key_metadata,
+            commit_uuid,
+            snapshot_properties,
+        )
+        .unwrap();
+
+        if use_starting_sequence_number {
+            if let Some(snapshot) = snapshot_produce_action
+                .tx
+                .base_table
+                .metadata()
+                .current_snapshot()
+            {
+                snapshot_produce_action
+                    .set_new_data_file_sequence_number(snapshot.sequence_number());
+            }
+        }
+
         Ok(Self {
-            snapshot_produce_action: SnapshotProduceAction::new(
-                tx,
-                snapshot_id,
-                key_metadata,
-                commit_uuid,
-                snapshot_properties,
-            )
-            .unwrap(),
+            snapshot_produce_action,
             target_size_bytes,
             min_count_to_merge,
             merge_enabled,
@@ -92,18 +120,18 @@ impl<'a> RewriteFilesAction<'a> {
     /// Add data files to the snapshot.
 
     pub fn add_data_files(
-        &mut self,
+        mut self,
         data_files: impl IntoIterator<Item = DataFile>,
-    ) -> Result<&mut Self> {
+    ) -> Result<Self> {
         self.snapshot_produce_action.add_data_files(data_files)?;
         Ok(self)
     }
 
     /// Add remove files to the snapshot.
     pub fn delete_files(
-        &mut self,
+        mut self,
         remove_data_files: impl IntoIterator<Item = DataFile>,
-    ) -> Result<&mut Self> {
+    ) -> Result<Self> {
         self.snapshot_produce_action
             .delete_files(remove_data_files)?;
         Ok(self)
@@ -123,6 +151,13 @@ impl<'a> RewriteFilesAction<'a> {
                 .await
         }
     }
+
+    pub fn new_data_file_sequence_number(mut self, seq: i64) -> Result<Self> {
+        self.snapshot_produce_action
+            .set_new_data_file_sequence_number(seq);
+
+        Ok(self)
+    }
 }
 
 impl SnapshotProduceOperation for RewriteFilesOperation {

crates/iceberg/src/transaction/snapshot.rs

@@ -89,6 +89,8 @@ pub(crate) struct SnapshotProduceAction<'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> SnapshotProduceAction<'a> {
@@ -112,6 +114,7 @@ impl<'a> SnapshotProduceAction<'a> {
             removed_delete_file_paths: HashSet::new(),
             manifest_counter: (0..),
             key_metadata,
+            new_data_file_sequence_number: None,
         })
     }
 
@@ -252,6 +255,7 @@ impl<'a> SnapshotProduceAction<'a> {
     async fn write_added_manifest(
         &mut self,
         added_data_files: Vec<DataFile>,
+        data_seq: Option<i64>,
     ) -> Result<ManifestFile> {
         let snapshot_id = self.snapshot_id;
         let format_version = self.tx.current_table.metadata().format_version();
@@ -279,7 +283,9 @@ impl<'a> SnapshotProduceAction<'a> {
         let manifest_entries = added_data_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 {
@@ -362,13 +368,18 @@ impl<'a> SnapshotProduceAction<'a> {
         let mut manifest_files = vec![];
         let data_files = std::mem::take(&mut self.added_data_files);
         let added_delete_files = std::mem::take(&mut self.added_delete_files);
+
         if !data_files.is_empty() {
-            let added_manifest = self.write_added_manifest(data_files).await?;
+            let added_manifest = self
+                .write_added_manifest(data_files, self.new_data_file_sequence_number)
+                .await?;
             manifest_files.push(added_manifest);
         }
 
         if !added_delete_files.is_empty() {
-            let added_delete_manifest = self.write_added_manifest(added_delete_files).await?;
+            let added_delete_manifest = self
+                .write_added_manifest(added_delete_files, self.new_data_file_sequence_number)
+                .await?;
             manifest_files.push(added_delete_manifest);
         }
 
@@ -479,6 +490,10 @@ impl<'a> SnapshotProduceAction<'a> {
         )?;
         Ok(self.tx)
     }
+
+    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 {

crates/integration_tests/tests/shared_tests/rewrite_files_test.rs

@@ -165,11 +165,13 @@ async fn test_rewrite_data_files() {
 
     // commit result again
     let tx = Transaction::new(&table);
-    let mut rewrite_action = tx.rewrite_files(None, vec![]).unwrap();
-    rewrite_action
+    let rewrite_action = tx
+        .rewrite_files(None, vec![])
+        .unwrap()
         .add_data_files(data_file_rewrite.clone())
+        .unwrap()
+        .delete_files(data_file.clone())
         .unwrap();
-    rewrite_action.delete_files(data_file.clone()).unwrap();
 
     let tx = rewrite_action.apply().await.unwrap();
     let table = tx.commit(&rest_catalog).await.unwrap();
@@ -279,9 +281,13 @@ async fn test_multiple_file_rewrite() {
     let data_file2 = data_file_writer.close().await.unwrap();
 
     let tx = Transaction::new(&table);
-    let mut rewrite_action = tx.rewrite_files(None, vec![]).unwrap();
-    rewrite_action.add_data_files(data_file1.clone()).unwrap();
-    rewrite_action.add_data_files(data_file2.clone()).unwrap();
+    let rewrite_action = tx
+        .rewrite_files(None, vec![])
+        .unwrap()
+        .add_data_files(data_file1.clone())
+        .unwrap()
+        .add_data_files(data_file2.clone())
+        .unwrap();
     let tx = rewrite_action.apply().await.unwrap();
     let table = tx.commit(&rest_catalog).await.unwrap();
 
@@ -357,10 +363,98 @@ async fn test_rewrite_nonexistent_file() {
     let nonexistent_data_file = valid_data_file.clone();
 
     let tx = Transaction::new(&table);
-    let mut rewrite_action = tx.rewrite_files(None, vec![]).unwrap();
+    let rewrite_action = tx.rewrite_files(None, vec![]).unwrap();
 
     // Attempt to delete the nonexistent file
     let result = rewrite_action.delete_files(nonexistent_data_file);
 
     assert!(result.is_ok());
 }
+
+#[tokio::test]
+async fn test_sequence_number_in_manifest_entry() {
+    let fixture = get_shared_containers();
+    let rest_catalog = RestCatalog::new(fixture.catalog_config.clone());
+    let ns = random_ns().await;
+    let schema = test_schema();
+
+    let table_creation = TableCreation::builder()
+        .name("t3".to_string())
+        .schema(schema.clone())
+        .build();
+
+    let table = rest_catalog
+        .create_table(ns.name(), table_creation)
+        .await
+        .unwrap();
+
+    let schema: Arc<arrow_schema::Schema> = Arc::new(
+        table
+            .metadata()
+            .current_schema()
+            .as_ref()
+            .try_into()
+            .unwrap(),
+    );
+    let location_generator = DefaultLocationGenerator::new(table.metadata().clone()).unwrap();
+    let file_name_generator = DefaultFileNameGenerator::new(
+        "test".to_string(),
+        None,
+        iceberg::spec::DataFileFormat::Parquet,
+    );
+    let parquet_writer_builder = ParquetWriterBuilder::new(
+        WriterProperties::default(),
+        table.metadata().current_schema().clone(),
+        table.file_io().clone(),
+        location_generator.clone(),
+        file_name_generator.clone(),
+    );
+    let data_file_writer_builder = DataFileWriterBuilder::new(parquet_writer_builder, None, 0);
+    let mut data_file_writer = data_file_writer_builder.clone().build().await.unwrap();
+    let col1 = StringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]);
+    let col2 = Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4)]);
+    let col3 = BooleanArray::from(vec![Some(true), Some(false), None, Some(false)]);
+    let batch = RecordBatch::try_new(schema.clone(), vec![
+        Arc::new(col1) as ArrayRef,
+        Arc::new(col2) as ArrayRef,
+        Arc::new(col3) as ArrayRef,
+    ])
+    .unwrap();
+    data_file_writer.write(batch.clone()).await.unwrap();
+    let data_file1 = data_file_writer.close().await.unwrap();
+
+    let mut data_file_writer = data_file_writer_builder.clone().build().await.unwrap();
+    data_file_writer.write(batch.clone()).await.unwrap();
+    let data_file2 = data_file_writer.close().await.unwrap();
+
+    // Commit with sequence number
+
+    let tx = Transaction::new(&table);
+    let rewrite_action = tx
+        .rewrite_files(None, vec![])
+        .unwrap()
+        .add_data_files(data_file1.clone())
+        .unwrap()
+        .add_data_files(data_file2.clone())
+        .unwrap();
+    // Set sequence number to 12345
+    let rewrite_action = rewrite_action.new_data_file_sequence_number(12345).unwrap();
+    let tx = rewrite_action.apply().await.unwrap();
+    let table = tx.commit(&rest_catalog).await.unwrap();
+
+    // Verify manifest entry has correct sequence number
+    let snapshot = table.metadata().current_snapshot().unwrap();
+    let manifest_list = snapshot
+        .load_manifest_list(table.file_io(), table.metadata())
+        .await
+        .unwrap();
+
+    assert_eq!(manifest_list.entries().len(), 1);
+
+    for manifest_file in manifest_list.entries() {
+        let manifest = manifest_file.load_manifest(table.file_io()).await.unwrap();
+        for entry in manifest.entries() {
+            assert_eq!(entry.sequence_number(), Some(12345));
+        }
+    }
+}