From 7d694b4af5658cd19e93aca26980b407a588248c Mon Sep 17 00:00:00 2001
From: Li0k <yuli@singularity-data.com>
Date: Mon, 10 Mar 2025 17:20:14 +0800
Subject: [PATCH] feat(iceberg): remove schemas

---
 crates/iceberg/src/catalog/mod.rs             | 22 ++++++++++++
 .../src/spec/table_metadata_builder.rs        | 35 +++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/crates/iceberg/src/catalog/mod.rs b/crates/iceberg/src/catalog/mod.rs
index 98b24b9256..e880641d7b 100644
--- a/crates/iceberg/src/catalog/mod.rs
+++ b/crates/iceberg/src/catalog/mod.rs
@@ -482,6 +482,12 @@ pub enum TableUpdate {
         /// Snapshot id to remove partition statistics for.
         snapshot_id: i64,
     },
+    /// Remove schemas
+    #[serde(rename_all = "kebab-case")]
+    RemoveSchemas {
+        /// Schema IDs to remove.
+        schema_ids: Vec<i32>,
+    },
 }
 
 impl TableUpdate {
@@ -525,6 +531,7 @@ impl TableUpdate {
             TableUpdate::RemovePartitionStatistics { snapshot_id } => {
                 Ok(builder.remove_partition_statistics(snapshot_id))
             }
+            TableUpdate::RemoveSchemas { schema_ids } => builder.remove_schemas(&schema_ids),
         }
     }
 }
@@ -2047,4 +2054,19 @@ mod tests {
             },
         )
     }
+
+    #[test]
+    fn test_remove_schema_update() {
+        test_serde_json(
+            r#"
+{
+    "action": "remove-schemas",
+    "schema-ids": [1, 2]
+}        
+        "#,
+            TableUpdate::RemoveSchemas {
+                schema_ids: vec![1, 2],
+            },
+        );
+    }
 }
diff --git a/crates/iceberg/src/spec/table_metadata_builder.rs b/crates/iceberg/src/spec/table_metadata_builder.rs
index 68ed76987b..eff732f856 100644
--- a/crates/iceberg/src/spec/table_metadata_builder.rs
+++ b/crates/iceberg/src/spec/table_metadata_builder.rs
@@ -1210,6 +1210,35 @@ impl TableMetadataBuilder {
     fn highest_sort_order_id(&self) -> Option<i64> {
         self.metadata.sort_orders.keys().max().copied()
     }
+
+    /// Remove schemas by their ids from the table metadata.
+    /// Does nothing if a schema id is not present. Active schemas should not be removed.
+    pub fn remove_schemas(mut self, schema_id_to_remove: &[i32]) -> Result<Self> {
+        if schema_id_to_remove.contains(&self.metadata.current_schema_id) {
+            return Err(Error::new(
+                ErrorKind::DataInvalid,
+                "Cannot remove current schema",
+            ));
+        }
+
+        if !schema_id_to_remove.is_empty() {
+            let mut removed_schemas = Vec::with_capacity(schema_id_to_remove.len());
+            self.metadata.schemas.retain(|id, _schema| {
+                if schema_id_to_remove.contains(id) {
+                    removed_schemas.push(*id);
+                    false
+                } else {
+                    true
+                }
+            });
+
+            self.changes.push(TableUpdate::RemoveSchemas {
+                schema_ids: removed_schemas,
+            });
+        }
+
+        Ok(self)
+    }
 }
 
 impl From<TableMetadataBuildResult> for TableMetadata {
@@ -2412,4 +2441,10 @@ mod tests {
             table.metadata().current_snapshot_id().unwrap()
         );
     }
+
+    #[test]
+    fn test_active_schema_cannot_be_removed() {
+        let builder = builder_without_changes(FormatVersion::V2);
+        builder.remove_schemas(&[0]).unwrap_err();
+    }
 }