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;
@@ -32,8 +33,7 @@ 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, TableMetadata, Transform,
-    MAIN_BRANCH,
+    SortDirection, SortField, SortOrder, Struct, StructType, Summary, Transform, MAIN_BRANCH,
 };
 use crate::table::Table;
 use crate::writer::file_writer::ParquetWriter;
@@ -45,7 +45,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 +55,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 +79,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 +107,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 +146,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)
         {
@@ -247,7 +249,8 @@ impl<'a> FastAppendAction<'a> {
         if !self
             .snapshot_produce_action
             .tx
-            .current_metadata
+            .current_table
+            .metadata()
             .default_spec
             .is_unpartitioned()
         {
@@ -258,9 +261,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?;
 
@@ -283,7 +286,7 @@ impl<'a> FastAppendAction<'a> {
             let mut manifest_stream = self
                 .snapshot_produce_action
                 .tx
-                .base_table
+                .current_table
                 .inspect()
                 .manifests()
                 .scan()
@@ -345,14 +348,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?;
 
@@ -470,7 +478,7 @@ impl<'a> SnapshotProduceAction<'a> {
             }
             Self::validate_partition_value(
                 data_file.partition(),
-                self.tx.current_metadata.default_partition_type(),
+                self.tx.current_table.metadata().default_partition_type(),
             )?;
         }
         self.added_data_files.extend(data_files);
@@ -480,20 +488,23 @@ 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.
     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 format_version = self.tx.current_table.metadata().format_version();
         let manifest_entries = added_data_files.into_iter().map(|data_file| {
             let builder = ManifestEntry::builder()
                 .status(crate::spec::ManifestStatus::Added)
@@ -511,14 +522,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 {
                 builder.build_v2_data()
@@ -558,7 +570,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,
@@ -576,28 +588,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,
             ),
         };
@@ -608,10 +620,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();
 
@@ -630,11 +642,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(),
                 },
             ],
         )?;
@@ -674,10 +686,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,
             },
         ];
 
@@ -693,7 +715,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(|| {