feat: support write to branch (#62)

* feat: support to branch

* fix: fix ref name

* fix: current_snapshot_id

* refactor: refactor interface

* fmt

Author: Li0k

crates/iceberg/src/transaction/append.rs

@@ -68,6 +68,11 @@ impl<'a> FastAppendAction<'a> {
         self
     }
 
+    pub fn with_to_branch(mut self, to_branch: String) -> Self {
+        self.snapshot_produce_action.set_target_branch(to_branch);
+        self
+    }
+
     /// Add data files to the snapshot.
     pub fn add_data_files(
         &mut self,
@@ -262,6 +267,11 @@ impl<'a> MergeAppendAction<'a> {
         })
     }
 
+    pub fn with_to_branch(mut self, to_branch: String) -> Result<MergeAppendAction<'a>> {
+        self.snapshot_produce_action.set_target_branch(to_branch);
+        Ok(self)
+    }
+
     /// Add data files to the snapshot.
     pub fn add_data_files(
         &mut self,
@@ -384,6 +394,94 @@ mod tests {
         assert_eq!(data_file, *manifest.entries()[0].data_file());
     }
 
+    #[tokio::test]
+    async fn test_fast_append_with_branch() {
+        let table = make_v2_minimal_table();
+        let tx = Transaction::new(&table);
+
+        // Test creating new branch
+        let branch_name = "test_branch";
+        let mut action = tx
+            .fast_append(None, None, vec![])
+            .unwrap()
+            .with_to_branch(branch_name.to_string());
+
+        let data_file = DataFileBuilder::default()
+            .partition_spec_id(0)
+            .content(DataContentType::Data)
+            .file_path("test/3.parquet".to_string())
+            .file_format(DataFileFormat::Parquet)
+            .file_size_in_bytes(100)
+            .record_count(1)
+            .partition_spec_id(table.metadata().default_partition_spec_id())
+            .partition(Struct::from_iter([Some(Literal::long(300))]))
+            .build()
+            .unwrap();
+        action.add_data_files(vec![data_file.clone()]).unwrap();
+        let tx = action.apply().await.unwrap();
+
+        // Check branch reference was created
+        assert!(
+            matches!(&tx.updates[1], TableUpdate::SetSnapshotRef { ref_name, .. }
+                if ref_name == branch_name)
+        );
+
+        // Test updating existing branch
+        let tx2 = Transaction::new(&table);
+        let mut action2 = tx2
+            .fast_append(None, None, vec![])
+            .unwrap()
+            .with_to_branch(branch_name.to_string());
+        action2.add_data_files(vec![data_file.clone()]).unwrap();
+        let tx2 = action2.apply().await.unwrap();
+
+        // Check requirements contain branch validation
+        assert!(tx2.requirements.iter().any(
+            |r| matches!(r, TableRequirement::RefSnapshotIdMatch { r#ref, .. }
+                if r#ref == branch_name)
+        ));
+    }
+
+    #[tokio::test]
+    async fn test_branch_operations() {
+        let table = make_v2_minimal_table();
+
+        // Test creating new branch
+        let branch_name = "test_branch";
+        let tx = Transaction::new(&table);
+        let mut action = tx
+            .fast_append(None, None, vec![])
+            .unwrap()
+            .with_to_branch(branch_name.to_string());
+
+        let data_file = DataFileBuilder::default()
+            .partition_spec_id(0)
+            .content(DataContentType::Data)
+            .file_path("test/3.parquet".to_string())
+            .file_format(DataFileFormat::Parquet)
+            .file_size_in_bytes(100)
+            .record_count(1)
+            .partition_spec_id(table.metadata().default_partition_spec_id())
+            .partition(Struct::from_iter([Some(Literal::long(300))]))
+            .build()
+            .unwrap();
+
+        action.add_data_files(vec![data_file.clone()]).unwrap();
+        let tx = action.apply().await.unwrap();
+
+        // Verify branch was created
+        assert!(matches!(
+            &tx.updates[1],
+            TableUpdate::SetSnapshotRef { ref_name, .. } if ref_name == branch_name
+        ));
+
+        // Verify requirements
+        assert!(tx.requirements.iter().any(
+            |r| matches!(r, TableRequirement::RefSnapshotIdMatch { r#ref, .. }
+                if r#ref == branch_name)
+        ));
+    }
+
     #[tokio::test]
     async fn test_add_existing_parquet_files_to_unpartitioned_table() {
         let mut fixture = TableTestFixture::new_unpartitioned();

crates/iceberg/src/transaction/rewrite_files.rs

@@ -118,7 +118,6 @@ impl<'a> RewriteFilesAction<'a> {
     }
 
     /// Add data files to the snapshot.
-
     pub fn add_data_files(
         mut self,
         data_files: impl IntoIterator<Item = DataFile>,
@@ -137,6 +136,11 @@ impl<'a> RewriteFilesAction<'a> {
         Ok(self)
     }
 
+    pub fn with_to_branch(mut self, to_branch: String) -> Result<Self> {
+        self.snapshot_produce_action.set_target_branch(to_branch);
+        Ok(self)
+    }
+
     /// Finished building the action and apply it to the transaction.
     pub async fn apply(self) -> Result<Transaction<'a>> {
         if self.merge_enabled {

crates/iceberg/src/transaction/snapshot.rs

@@ -91,6 +91,8 @@ pub(crate) struct SnapshotProduceAction<'a> {
     manifest_counter: RangeFrom<u64>,
 
     new_data_file_sequence_number: Option<i64>,
+
+    target_branch: String,
 }
 
 impl<'a> SnapshotProduceAction<'a> {
@@ -115,6 +117,7 @@ impl<'a> SnapshotProduceAction<'a> {
             manifest_counter: (0..),
             key_metadata,
             new_data_file_sequence_number: None,
+            target_branch: MAIN_BRANCH.to_string(),
         })
     }
 
@@ -426,28 +429,39 @@ impl<'a> SnapshotProduceAction<'a> {
         let new_manifests = self
             .manifest_file(&snapshot_produce_operation, &process)
             .await?;
+
         let next_seq_num = self.tx.current_table.metadata().next_sequence_number();
 
         let summary = self.summary(&snapshot_produce_operation);
 
         let manifest_list_path = self.generate_manifest_list_file_path(0);
 
+        let parent_snapshot = self
+            .tx
+            .current_table
+            .metadata()
+            .snapshot_for_ref(&self.target_branch);
+
+        let parent_snapshot_id = parent_snapshot
+            .map(|s| Some(s.snapshot_id()))
+            .unwrap_or(None);
+
         let mut manifest_list_writer = match self.tx.current_table.metadata().format_version() {
             FormatVersion::V1 => ManifestListWriter::v1(
                 self.tx
                     .current_table
                     .file_io()
                     .new_output(manifest_list_path.clone())?,
                 self.snapshot_id,
-                self.tx.current_table.metadata().current_snapshot_id(),
+                parent_snapshot_id,
             ),
             FormatVersion::V2 => ManifestListWriter::v2(
                 self.tx
                     .current_table
                     .file_io()
                     .new_output(manifest_list_path.clone())?,
                 self.snapshot_id,
-                self.tx.current_table.metadata().current_snapshot_id(),
+                parent_snapshot_id,
                 next_seq_num,
             ),
         };
@@ -458,7 +472,7 @@ impl<'a> SnapshotProduceAction<'a> {
         let new_snapshot = Snapshot::builder()
             .with_manifest_list(manifest_list_path)
             .with_snapshot_id(self.snapshot_id)
-            .with_parent_snapshot_id(self.tx.current_table.metadata().current_snapshot_id())
+            .with_parent_snapshot_id(parent_snapshot_id)
             .with_sequence_number(next_seq_num)
             .with_summary(summary)
             .with_schema_id(self.tx.current_table.metadata().current_schema_id())
@@ -471,7 +485,7 @@ impl<'a> SnapshotProduceAction<'a> {
                     snapshot: new_snapshot,
                 },
                 TableUpdate::SetSnapshotRef {
-                    ref_name: MAIN_BRANCH.to_string(),
+                    ref_name: self.target_branch.clone(),
                     reference: SnapshotReference::new(
                         self.snapshot_id,
                         SnapshotRetention::branch(None, None, None),
@@ -483,8 +497,8 @@ impl<'a> SnapshotProduceAction<'a> {
                     uuid: self.tx.current_table.metadata().uuid(),
                 },
                 TableRequirement::RefSnapshotIdMatch {
-                    r#ref: MAIN_BRANCH.to_string(),
-                    snapshot_id: self.tx.current_table.metadata().current_snapshot_id(),
+                    r#ref: self.target_branch.clone(),
+                    snapshot_id: parent_snapshot_id,
                 },
             ],
         )?;
@@ -494,6 +508,10 @@ impl<'a> SnapshotProduceAction<'a> {
     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 fn set_target_branch(&mut self, target_branch: String) {
+        self.target_branch = target_branch;
+    }
 }
 
 pub(crate) struct MergeManifestProcess {
@@ -519,6 +537,7 @@ impl ManifestProcess for MergeManifestProcess {
         let (unmerge_data_manifest, unmerge_delete_manifest): (Vec<_>, Vec<_>) = manifests
             .into_iter()
             .partition(|manifest| matches!(manifest.content, ManifestContentType::Data));
+
         let mut data_manifest = {
             let manifest_merge_manager = MergeManifestManager::new(
                 self.target_size_bytes,
@@ -529,6 +548,7 @@ impl ManifestProcess for MergeManifestProcess {
                 .merge_manifest(snapshot_produce, unmerge_data_manifest)
                 .await?
         };
+
         data_manifest.extend(unmerge_delete_manifest);
         Ok(data_manifest)
     }