add apply in transaction to support stack action

Author: ZENOTME

crates/iceberg/src/transaction.rs

@@ -30,7 +30,8 @@ use crate::io::OutputFile;
 use crate::spec::{
     DataFile, DataFileFormat, FormatVersion, ManifestEntry, ManifestFile, ManifestListWriter,
     ManifestWriterBuilder, NullOrder, Operation, Snapshot, SnapshotReference, SnapshotRetention,
-    SortDirection, SortField, SortOrder, Struct, StructType, Summary, Transform, MAIN_BRANCH,
+    SortDirection, SortField, SortOrder, Struct, StructType, Summary, TableMetadata, Transform,
+    MAIN_BRANCH,
 };
 use crate::table::Table;
 use crate::TableUpdate::UpgradeFormatVersion;
@@ -40,7 +41,8 @@ const META_ROOT_PATH: &str = "metadata";
 
 /// Table transaction.
 pub struct Transaction<'a> {
-    table: &'a Table,
+    base_table: &'a Table,
+    current_metadata: TableMetadata,
     updates: Vec<TableUpdate>,
     requirements: Vec<TableRequirement>,
 }
@@ -49,38 +51,59 @@ impl<'a> Transaction<'a> {
     /// Creates a new transaction.
     pub fn new(table: &'a Table) -> Self {
         Self {
-            table,
+            base_table: table,
+            current_metadata: table.metadata().clone(),
             updates: vec![],
             requirements: vec![],
         }
     }
 
-    fn append_updates(&mut self, updates: Vec<TableUpdate>) -> Result<()> {
-        for update in &updates {
-            for up in &self.updates {
-                if discriminant(up) == discriminant(update) {
-                    return Err(Error::new(
-                        ErrorKind::DataInvalid,
-                        format!(
-                            "Cannot apply update with same type at same time: {:?}",
-                            update
-                        ),
-                    ));
-                }
-            }
+    fn update_table_metadata(&mut self, updates: &[TableUpdate]) -> Result<()> {
+        let mut metadata_builder = self.current_metadata.clone().into_builder(None);
+        for update in updates {
+            metadata_builder = update.clone().apply(metadata_builder)?;
         }
-        self.updates.extend(updates);
+
+        self.current_metadata = metadata_builder.build()?.metadata;
+
         Ok(())
     }
 
-    fn append_requirements(&mut self, requirements: Vec<TableRequirement>) -> Result<()> {
-        self.requirements.extend(requirements);
+    fn apply(
+        &mut self,
+        updates: Vec<TableUpdate>,
+        requirements: Vec<TableRequirement>,
+    ) -> Result<()> {
+        for requirment in &requirements {
+            requirment.check(Some(&self.current_metadata))?;
+        }
+
+        self.update_table_metadata(&updates)?;
+
+        self.updates.extend(updates);
+
+        // For the requirements, it does not make sense to add a requirement more than once
+        // For example, you cannot assert that the current schema has two different IDs
+        for new_requirment in requirements {
+            if self
+                .requirements
+                .iter()
+                .map(discriminant)
+                .all(|d| d != discriminant(&new_requirment))
+            {
+                self.requirements.push(new_requirment);
+            }
+        }
+
+        // # TODO
+        // Support auto commit later.
+
         Ok(())
     }
 
     /// Sets table to a new version.
     pub fn upgrade_table_version(mut self, format_version: FormatVersion) -> Result<Self> {
-        let current_version = self.table.metadata().format_version();
+        let current_version = self.current_metadata.format_version();
         match current_version.cmp(&format_version) {
             Ordering::Greater => {
                 return Err(Error::new(
@@ -92,7 +115,7 @@ impl<'a> Transaction<'a> {
                 ));
             }
             Ordering::Less => {
-                self.append_updates(vec![UpgradeFormatVersion { format_version }])?;
+                self.apply(vec![UpgradeFormatVersion { format_version }], vec![])?;
             }
             Ordering::Equal => {
                 // Do nothing.
@@ -103,7 +126,7 @@ impl<'a> Transaction<'a> {
 
     /// Update table's property.
     pub fn set_properties(mut self, props: HashMap<String, String>) -> Result<Self> {
-        self.append_updates(vec![TableUpdate::SetProperties { updates: props }])?;
+        self.apply(vec![TableUpdate::SetProperties { updates: props }], vec![])?;
         Ok(self)
     }
 
@@ -119,8 +142,7 @@ impl<'a> Transaction<'a> {
         };
         let mut snapshot_id = generate_random_id();
         while self
-            .table
-            .metadata()
+            .current_metadata
             .snapshots()
             .any(|s| s.snapshot_id() == snapshot_id)
         {
@@ -155,14 +177,17 @@ impl<'a> Transaction<'a> {
 
     /// Remove properties in table.
     pub fn remove_properties(mut self, keys: Vec<String>) -> Result<Self> {
-        self.append_updates(vec![TableUpdate::RemoveProperties { removals: keys }])?;
+        self.apply(
+            vec![TableUpdate::RemoveProperties { removals: keys }],
+            vec![],
+        )?;
         Ok(self)
     }
 
     /// Commit transaction.
     pub async fn commit(self, catalog: &impl Catalog) -> Result<Table> {
         let table_commit = TableCommit::builder()
-            .ident(self.table.identifier().clone())
+            .ident(self.base_table.identifier().clone())
             .updates(self.updates)
             .requirements(self.requirements)
             .build();
@@ -231,14 +256,14 @@ impl SnapshotProduceOperation for FastAppendOperation {
         &self,
         snapshot_produce: &SnapshotProduceAction<'_>,
     ) -> Result<Vec<ManifestFile>> {
-        let Some(snapshot) = snapshot_produce.tx.table.metadata().current_snapshot() else {
+        let Some(snapshot) = snapshot_produce.tx.current_metadata.current_snapshot() else {
             return Ok(vec![]);
         };
 
         let manifest_list = snapshot
             .load_manifest_list(
-                snapshot_produce.tx.table.file_io(),
-                &snapshot_produce.tx.table.metadata_ref(),
+                snapshot_produce.tx.base_table.file_io(),
+                &snapshot_produce.tx.current_metadata,
             )
             .await?;
 
@@ -355,7 +380,7 @@ impl<'a> SnapshotProduceAction<'a> {
             }
             Self::validate_partition_value(
                 data_file.partition(),
-                self.tx.table.metadata().default_partition_type(),
+                self.tx.current_metadata.default_partition_type(),
             )?;
         }
         self.added_data_files.extend(data_files);
@@ -365,24 +390,25 @@ impl<'a> SnapshotProduceAction<'a> {
     fn new_manifest_output(&mut self) -> Result<OutputFile> {
         let new_manifest_path = format!(
             "{}/{}/{}-m{}.{}",
-            self.tx.table.metadata().location(),
+            self.tx.current_metadata.location(),
             META_ROOT_PATH,
             self.commit_uuid,
             self.manifest_counter.next().unwrap(),
             DataFileFormat::Avro
         );
-        self.tx.table.file_io().new_output(new_manifest_path)
+        self.tx.base_table.file_io().new_output(new_manifest_path)
     }
 
     // Write manifest file for added data files and return the ManifestFile for ManifestList.
     async fn write_added_manifest(&mut self) -> Result<ManifestFile> {
         let added_data_files = std::mem::take(&mut self.added_data_files);
         let snapshot_id = self.snapshot_id;
+        let format_version = self.tx.current_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 self.tx.table.metadata().format_version() == FormatVersion::V1 {
+            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
@@ -395,15 +421,14 @@ impl<'a> SnapshotProduceAction<'a> {
                 self.new_manifest_output()?,
                 Some(self.snapshot_id),
                 self.key_metadata.clone(),
-                self.tx.table.metadata().current_schema().clone(),
+                self.tx.current_metadata.current_schema().clone(),
                 self.tx
-                    .table
-                    .metadata()
+                    .current_metadata
                     .default_partition_spec()
                     .as_ref()
                     .clone(),
             );
-            if self.tx.table.metadata().format_version() == FormatVersion::V1 {
+            if self.tx.current_metadata.format_version() == FormatVersion::V1 {
                 builder.build_v1()
             } else {
                 builder.build_v2_data()
@@ -443,7 +468,7 @@ impl<'a> SnapshotProduceAction<'a> {
     fn generate_manifest_list_file_path(&self, attempt: i64) -> String {
         format!(
             "{}/{}/snap-{}-{}-{}.{}",
-            self.tx.table.metadata().location(),
+            self.tx.current_metadata.location(),
             META_ROOT_PATH,
             self.snapshot_id,
             attempt,
@@ -461,28 +486,28 @@ impl<'a> SnapshotProduceAction<'a> {
         let new_manifests = self
             .manifest_file(&snapshot_produce_operation, &process)
             .await?;
-        let next_seq_num = self.tx.table.metadata().next_sequence_number();
+        let next_seq_num = self.tx.current_metadata.next_sequence_number();
 
         let summary = self.summary(&snapshot_produce_operation);
 
         let manifest_list_path = self.generate_manifest_list_file_path(0);
 
-        let mut manifest_list_writer = match self.tx.table.metadata().format_version() {
+        let mut manifest_list_writer = match self.tx.current_metadata.format_version() {
             FormatVersion::V1 => ManifestListWriter::v1(
                 self.tx
-                    .table
+                    .base_table
                     .file_io()
                     .new_output(manifest_list_path.clone())?,
                 self.snapshot_id,
-                self.tx.table.metadata().current_snapshot_id(),
+                self.tx.current_metadata.current_snapshot_id(),
             ),
             FormatVersion::V2 => ManifestListWriter::v2(
                 self.tx
-                    .table
+                    .base_table
                     .file_io()
                     .new_output(manifest_list_path.clone())?,
                 self.snapshot_id,
-                self.tx.table.metadata().current_snapshot_id(),
+                self.tx.current_metadata.current_snapshot_id(),
                 next_seq_num,
             ),
         };
@@ -493,34 +518,36 @@ 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.table.metadata().current_snapshot_id())
+            .with_parent_snapshot_id(self.tx.current_metadata.current_snapshot_id())
             .with_sequence_number(next_seq_num)
             .with_summary(summary)
-            .with_schema_id(self.tx.table.metadata().current_schema_id())
+            .with_schema_id(self.tx.current_metadata.current_schema_id())
             .with_timestamp_ms(commit_ts)
             .build();
 
-        self.tx.append_updates(vec![
-            TableUpdate::AddSnapshot {
-                snapshot: new_snapshot,
-            },
-            TableUpdate::SetSnapshotRef {
-                ref_name: MAIN_BRANCH.to_string(),
-                reference: SnapshotReference::new(
-                    self.snapshot_id,
-                    SnapshotRetention::branch(None, None, None),
-                ),
-            },
-        ])?;
-        self.tx.append_requirements(vec![
-            TableRequirement::UuidMatch {
-                uuid: self.tx.table.metadata().uuid(),
-            },
-            TableRequirement::RefSnapshotIdMatch {
-                r#ref: MAIN_BRANCH.to_string(),
-                snapshot_id: self.tx.table.metadata().current_snapshot_id(),
-            },
-        ])?;
+        self.tx.apply(
+            vec![
+                TableUpdate::AddSnapshot {
+                    snapshot: new_snapshot.clone(),
+                },
+                TableUpdate::SetSnapshotRef {
+                    ref_name: MAIN_BRANCH.to_string(),
+                    reference: SnapshotReference::new(
+                        self.snapshot_id,
+                        SnapshotRetention::branch(None, None, None),
+                    ),
+                },
+            ],
+            vec![
+                TableRequirement::UuidMatch {
+                    uuid: self.tx.current_metadata.uuid(),
+                },
+                TableRequirement::RefSnapshotIdMatch {
+                    r#ref: MAIN_BRANCH.to_string(),
+                    snapshot_id: self.tx.current_metadata.current_snapshot_id(),
+                },
+            ],
+        )?;
         Ok(self.tx)
     }
 }
@@ -557,15 +584,14 @@ impl<'a> ReplaceSortOrderAction<'a> {
 
         let requirements = vec![
             TableRequirement::CurrentSchemaIdMatch {
-                current_schema_id: self.tx.table.metadata().current_schema().schema_id(),
+                current_schema_id: self.tx.current_metadata.current_schema().schema_id(),
             },
             TableRequirement::DefaultSortOrderIdMatch {
-                default_sort_order_id: self.tx.table.metadata().default_sort_order().order_id,
+                default_sort_order_id: self.tx.current_metadata.default_sort_order().order_id,
             },
         ];
 
-        self.tx.append_requirements(requirements)?;
-        self.tx.append_updates(updates)?;
+        self.tx.apply(updates, requirements)?;
         Ok(self.tx)
     }
 
@@ -577,8 +603,7 @@ impl<'a> ReplaceSortOrderAction<'a> {
     ) -> Result<Self> {
         let field_id = self
             .tx
-            .table
-            .metadata()
+            .current_metadata
             .current_schema()
             .field_id_by_name(name)
             .ok_or_else(|| {
@@ -806,14 +831,15 @@ mod tests {
         assert!(
             matches!((&tx.updates[0],&tx.updates[1]), (TableUpdate::AddSnapshot { snapshot },TableUpdate::SetSnapshotRef { reference,ref_name }) if snapshot.snapshot_id() == reference.snapshot_id && ref_name == MAIN_BRANCH)
         );
+        // requriments is based on original table metadata
         assert_eq!(
             vec![
                 TableRequirement::UuidMatch {
-                    uuid: tx.table.metadata().uuid()
+                    uuid: table.metadata().uuid()
                 },
                 TableRequirement::RefSnapshotIdMatch {
                     r#ref: MAIN_BRANCH.to_string(),
-                    snapshot_id: tx.table.metadata().current_snapshot_id
+                    snapshot_id: table.metadata().current_snapshot_id()
                 }
             ],
             tx.requirements
@@ -853,20 +879,4 @@ mod tests {
         );
         assert_eq!(data_file, *manifest.entries()[0].data_file());
     }
-
-    #[test]
-    fn test_do_same_update_in_same_transaction() {
-        let table = make_v2_table();
-        let tx = Transaction::new(&table);
-        let tx = tx
-            .remove_properties(vec!["a".to_string(), "b".to_string()])
-            .unwrap();
-
-        let tx = tx.remove_properties(vec!["c".to_string(), "d".to_string()]);
-
-        assert!(
-            tx.is_err(),
-            "Should not allow to do same kinds update in same transaction"
-        );
-    }
 }