(iceberg-rust) branch main updated: feat: Add RemovePartitionSpecs table update (#804)

[xu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

xuanwo pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-rust.git


The following commit(s) were added to refs/heads/main by this push:
     new fa8c9485 feat: Add RemovePartitionSpecs table update (#804)
fa8c9485 is described below

commit fa8c948538ebb7792365a295cc0ffa9b7fd0a127
Author: Christian <ch...@hansetag.com>
AuthorDate: Tue Dec 17 02:40:27 2024 +0100

    feat: Add RemovePartitionSpecs table update (#804)
    
    * RemovePartitionSpecs update
    
    * address comments
    
    * fix comments
---
 crates/iceberg/src/catalog/mod.rs                 | 24 +++++++++++
 crates/iceberg/src/spec/table_metadata_builder.rs | 52 +++++++++++++++++++++++
 2 files changed, 76 insertions(+)

diff --git a/crates/iceberg/src/catalog/mod.rs b/crates/iceberg/src/catalog/mod.rs
index ffafc66f..cbda6c90 100644
--- a/crates/iceberg/src/catalog/mod.rs
+++ b/crates/iceberg/src/catalog/mod.rs
@@ -440,6 +440,12 @@ pub enum TableUpdate {
         /// Properties to remove
         removals: Vec<String>,
     },
+    /// Remove partition specs
+    #[serde(rename_all = "kebab-case")]
+    RemovePartitionSpecs {
+        /// Partition spec ids to remove.
+        spec_ids: Vec<i32>,
+    },
     /// Set statistics for a snapshot
     #[serde(with = "_serde_set_statistics")]
     SetStatistics {
@@ -494,6 +500,9 @@ impl TableUpdate {
             TableUpdate::UpgradeFormatVersion { format_version } => {
                 builder.upgrade_format_version(format_version)
             }
+            TableUpdate::RemovePartitionSpecs { spec_ids } => {
+                builder.remove_partition_specs(&spec_ids)
+            }
             TableUpdate::SetStatistics { statistics } => Ok(builder.set_statistics(statistics)),
             TableUpdate::RemoveStatistics { snapshot_id } => {
                 Ok(builder.remove_statistics(snapshot_id))
@@ -1910,6 +1919,21 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_remove_partition_specs_update() {
+        test_serde_json(
+            r#"
+{
+    "action": "remove-partition-specs",
+    "spec-ids": [1, 2]
+}        
+        "#,
+            TableUpdate::RemovePartitionSpecs {
+                spec_ids: vec![1, 2],
+            },
+        );
+    }
+
     #[test]
     fn test_set_statistics_file() {
         test_serde_json(
diff --git a/crates/iceberg/src/spec/table_metadata_builder.rs b/crates/iceberg/src/spec/table_metadata_builder.rs
index 5b4ff123..4e8c1ae7 100644
--- a/crates/iceberg/src/spec/table_metadata_builder.rs
+++ b/crates/iceberg/src/spec/table_metadata_builder.rs
@@ -784,6 +784,36 @@ impl TableMetadataBuilder {
             .set_default_partition_spec(Self::LAST_ADDED)
     }
 
+    /// Remove partition specs by their ids from the table metadata.
+    /// Does nothing if a spec id is not present. Active partition specs
+    /// should not be removed.
+    ///
+    /// # Errors
+    /// - Cannot remove the default partition spec.
+    pub fn remove_partition_specs(mut self, spec_ids: &[i32]) -> Result<Self> {
+        if spec_ids.contains(&self.metadata.default_spec.spec_id()) {
+            return Err(Error::new(
+                ErrorKind::DataInvalid,
+                "Cannot remove default partition spec",
+            ));
+        }
+
+        let mut removed_specs = Vec::with_capacity(spec_ids.len());
+        spec_ids.iter().for_each(|id| {
+            if self.metadata.partition_specs.remove(id).is_some() {
+                removed_specs.push(*id);
+            }
+        });
+
+        if !removed_specs.is_empty() {
+            self.changes.push(TableUpdate::RemovePartitionSpecs {
+                spec_ids: removed_specs,
+            });
+        }
+
+        Ok(self)
+    }
+
     /// Add a sort order to the table metadata.
     ///
     /// The spec is bound eagerly to the current schema and must be valid for it.
@@ -1584,6 +1614,21 @@ mod tests {
         pretty_assertions::assert_eq!(build_result.changes[0], TableUpdate::AddSpec {
             spec: expected_change
         });
+
+        // Remove the spec
+        let build_result = build_result
+            .metadata
+            .into_builder(Some(
+                "s3://bucket/test/location/metadata/metadata1.json".to_string(),
+            ))
+            .remove_partition_specs(&[1])
+            .unwrap()
+            .build()
+            .unwrap();
+
+        assert_eq!(build_result.changes.len(), 1);
+        assert_eq!(build_result.metadata.partition_specs.len(), 1);
+        assert!(build_result.metadata.partition_spec_by_id(1).is_none());
     }
 
     #[test]
@@ -2193,6 +2238,13 @@ mod tests {
             .contains("Cannot add snapshot with sequence number"));
     }
 
+    #[test]
+    fn test_default_spec_cannot_be_removed() {
+        let builder = builder_without_changes(FormatVersion::V2);
+
+        builder.remove_partition_specs(&[0]).unwrap_err();
+    }
+
     #[test]
     fn test_statistics() {
         let builder = builder_without_changes(FormatVersion::V2);