store current table instead of current metadata

Author: ZENOTME

crates/iceberg/src/table.rs

@@ -162,6 +162,10 @@ pub struct Table {
 }
 
 impl Table {
+    pub(crate) fn with_metadata(&mut self, metadata: TableMetadataRef) {
+        self.metadata = metadata;
+    }
+
     /// Returns a TableBuilder to build a table
     pub fn builder() -> TableBuilder {
         TableBuilder::new()

crates/iceberg/src/transaction.rs

@@ -22,6 +22,7 @@ use std::collections::{HashMap, HashSet};
 use std::future::Future;
 use std::mem::discriminant;
 use std::ops::RangeFrom;
+use std::sync::Arc;
 
 use arrow_array::StringArray;
 use futures::TryStreamExt;
@@ -45,7 +46,7 @@ const META_ROOT_PATH: &str = "metadata";
 /// Table transaction.
 pub struct Transaction<'a> {
     base_table: &'a Table,
-    current_metadata: TableMetadata,
+    current_table: Table,
     updates: Vec<TableUpdate>,
     requirements: Vec<TableRequirement>,
 }
@@ -55,19 +56,20 @@ impl<'a> Transaction<'a> {
     pub fn new(table: &'a Table) -> Self {
         Self {
             base_table: table,
-            current_metadata: table.metadata().clone(),
+            current_table: table.clone(),
             updates: vec![],
             requirements: vec![],
         }
     }
 
     fn update_table_metadata(&mut self, updates: &[TableUpdate]) -> Result<()> {
-        let mut metadata_builder = self.current_metadata.clone().into_builder(None);
+        let mut metadata_builder = self.current_table.metadata().clone().into_builder(None);
         for update in updates {
             metadata_builder = update.clone().apply(metadata_builder)?;
         }
 
-        self.current_metadata = metadata_builder.build()?.metadata;
+        self.current_table
+            .with_metadata(Arc::new(metadata_builder.build()?.metadata));
 
         Ok(())
     }
@@ -78,7 +80,7 @@ impl<'a> Transaction<'a> {
         requirements: Vec<TableRequirement>,
     ) -> Result<()> {
         for requirement in &requirements {
-            requirement.check(Some(&self.current_metadata))?;
+            requirement.check(Some(self.current_table.metadata()))?;
         }
 
         self.update_table_metadata(&updates)?;
@@ -106,7 +108,7 @@ impl<'a> Transaction<'a> {
 
     /// Sets table to a new version.
     pub fn upgrade_table_version(mut self, format_version: FormatVersion) -> Result<Self> {
-        let current_version = self.current_metadata.format_version();
+        let current_version = self.current_table.metadata().format_version();
         match current_version.cmp(&format_version) {
             Ordering::Greater => {
                 return Err(Error::new(
@@ -145,7 +147,8 @@ impl<'a> Transaction<'a> {
         };
         let mut snapshot_id = generate_random_id();
         while self
-            .current_metadata
+            .current_table
+            .metadata()
             .snapshots()
             .any(|s| s.snapshot_id() == snapshot_id)
         {
@@ -239,7 +242,8 @@ impl<'a> FastAppendAction<'a> {
         if !self
             .snapshot_produce_action
             .tx
-            .current_metadata
+            .current_table
+            .metadata()
             .default_spec
             .is_unpartitioned()
         {
@@ -250,9 +254,9 @@ impl<'a> FastAppendAction<'a> {
         }
 
         let data_files = ParquetWriter::parquet_files_to_data_files(
-            self.snapshot_produce_action.tx.base_table.file_io(),
+            self.snapshot_produce_action.tx.current_table.file_io(),
             file_path,
-            &self.snapshot_produce_action.tx.current_metadata,
+            self.snapshot_produce_action.tx.current_table.metadata(),
         )
         .await?;
 
@@ -274,7 +278,7 @@ impl<'a> FastAppendAction<'a> {
         let mut manifest_stream = self
             .snapshot_produce_action
             .tx
-            .base_table
+            .current_table
             .inspect()
             .manifests()
             .scan()
@@ -335,14 +339,19 @@ impl SnapshotProduceOperation for FastAppendOperation {
         &self,
         snapshot_produce: &SnapshotProduceAction<'_>,
     ) -> Result<Vec<ManifestFile>> {
-        let Some(snapshot) = snapshot_produce.tx.current_metadata.current_snapshot() else {
+        let Some(snapshot) = snapshot_produce
+            .tx
+            .current_table
+            .metadata()
+            .current_snapshot()
+        else {
             return Ok(vec![]);
         };
 
         let manifest_list = snapshot
             .load_manifest_list(
-                snapshot_produce.tx.base_table.file_io(),
-                &snapshot_produce.tx.current_metadata,
+                snapshot_produce.tx.current_table.file_io(),
+                snapshot_produce.tx.current_table.metadata(),
             )
             .await?;
 
@@ -456,7 +465,7 @@ impl<'a> SnapshotProduceAction<'a> {
         for data_file in data_files {
             Self::validate_partition_value(
                 data_file.partition(),
-                self.tx.current_metadata.default_partition_type(),
+                self.tx.current_table.metadata().default_partition_type(),
             )?;
             if data_file.content_type() == DataContentType::Data {
                 self.added_data_files.push(data_file);
@@ -470,13 +479,16 @@ impl<'a> SnapshotProduceAction<'a> {
     fn new_manifest_output(&mut self) -> Result<OutputFile> {
         let new_manifest_path = format!(
             "{}/{}/{}-m{}.{}",
-            self.tx.current_metadata.location(),
+            self.tx.current_table.metadata().location(),
             META_ROOT_PATH,
             self.commit_uuid,
             self.manifest_counter.next().unwrap(),
             DataFileFormat::Avro
         );
-        self.tx.base_table.file_io().new_output(new_manifest_path)
+        self.tx
+            .current_table
+            .file_io()
+            .new_output(new_manifest_path)
     }
 
     // Write manifest file for added data files and return the ManifestFile for ManifestList.
@@ -485,7 +497,7 @@ impl<'a> SnapshotProduceAction<'a> {
         added_data_files: Vec<DataFile>,
     ) -> Result<ManifestFile> {
         let snapshot_id = self.snapshot_id;
-        let format_version = self.tx.current_metadata.format_version();
+        let format_version = self.tx.current_table.metadata().format_version();
         let content_type = {
             let mut data_num = 0;
             let mut delete_num = 0;
@@ -524,14 +536,15 @@ impl<'a> SnapshotProduceAction<'a> {
                 self.new_manifest_output()?,
                 Some(self.snapshot_id),
                 self.key_metadata.clone(),
-                self.tx.current_metadata.current_schema().clone(),
+                self.tx.current_table.metadata().current_schema().clone(),
                 self.tx
-                    .current_metadata
+                    .current_table
+                    .metadata()
                     .default_partition_spec()
                     .as_ref()
                     .clone(),
             );
-            if self.tx.current_metadata.format_version() == FormatVersion::V1 {
+            if self.tx.current_table.metadata().format_version() == FormatVersion::V1 {
                 builder.build_v1()
             } else {
                 match content_type {
@@ -581,7 +594,7 @@ impl<'a> SnapshotProduceAction<'a> {
     fn generate_manifest_list_file_path(&self, attempt: i64) -> String {
         format!(
             "{}/{}/snap-{}-{}-{}.{}",
-            self.tx.current_metadata.location(),
+            self.tx.current_table.metadata().location(),
             META_ROOT_PATH,
             self.snapshot_id,
             attempt,
@@ -599,28 +612,28 @@ impl<'a> SnapshotProduceAction<'a> {
         let new_manifests = self
             .manifest_file(&snapshot_produce_operation, &process)
             .await?;
-        let next_seq_num = self.tx.current_metadata.next_sequence_number();
+        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 mut manifest_list_writer = match self.tx.current_metadata.format_version() {
+        let mut manifest_list_writer = match self.tx.current_table.metadata().format_version() {
             FormatVersion::V1 => ManifestListWriter::v1(
                 self.tx
-                    .base_table
+                    .current_table
                     .file_io()
                     .new_output(manifest_list_path.clone())?,
                 self.snapshot_id,
-                self.tx.current_metadata.current_snapshot_id(),
+                self.tx.current_table.metadata().current_snapshot_id(),
             ),
             FormatVersion::V2 => ManifestListWriter::v2(
                 self.tx
-                    .base_table
+                    .current_table
                     .file_io()
                     .new_output(manifest_list_path.clone())?,
                 self.snapshot_id,
-                self.tx.current_metadata.current_snapshot_id(),
+                self.tx.current_table.metadata().current_snapshot_id(),
                 next_seq_num,
             ),
         };
@@ -631,10 +644,10 @@ 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_metadata.current_snapshot_id())
+            .with_parent_snapshot_id(self.tx.current_table.metadata().current_snapshot_id())
             .with_sequence_number(next_seq_num)
             .with_summary(summary)
-            .with_schema_id(self.tx.current_metadata.current_schema_id())
+            .with_schema_id(self.tx.current_table.metadata().current_schema_id())
             .with_timestamp_ms(commit_ts)
             .build();
 
@@ -653,11 +666,11 @@ impl<'a> SnapshotProduceAction<'a> {
             ],
             vec![
                 TableRequirement::UuidMatch {
-                    uuid: self.tx.current_metadata.uuid(),
+                    uuid: self.tx.current_table.metadata().uuid(),
                 },
                 TableRequirement::RefSnapshotIdMatch {
                     r#ref: MAIN_BRANCH.to_string(),
-                    snapshot_id: self.tx.current_metadata.current_snapshot_id(),
+                    snapshot_id: self.tx.current_table.metadata().current_snapshot_id(),
                 },
             ],
         )?;
@@ -697,10 +710,20 @@ impl<'a> ReplaceSortOrderAction<'a> {
 
         let requirements = vec![
             TableRequirement::CurrentSchemaIdMatch {
-                current_schema_id: self.tx.current_metadata.current_schema().schema_id(),
+                current_schema_id: self
+                    .tx
+                    .current_table
+                    .metadata()
+                    .current_schema()
+                    .schema_id(),
             },
             TableRequirement::DefaultSortOrderIdMatch {
-                default_sort_order_id: self.tx.current_metadata.default_sort_order().order_id,
+                default_sort_order_id: self
+                    .tx
+                    .current_table
+                    .metadata()
+                    .default_sort_order()
+                    .order_id,
             },
         ];
 
@@ -716,7 +739,8 @@ impl<'a> ReplaceSortOrderAction<'a> {
     ) -> Result<Self> {
         let field_id = self
             .tx
-            .current_metadata
+            .current_table
+            .metadata()
             .current_schema()
             .field_id_by_name(name)
             .ok_or_else(|| {