add unit test

Author: dentiny

crates/iceberg/src/transaction/update_schema.rs

@@ -15,20 +15,20 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::collections::{HashMap, HashSet};
+use std::collections::HashSet;
 use std::sync::Arc;
 
 use async_trait::async_trait;
-use uuid::Uuid;
 
 use crate::error::{Error, ErrorKind, Result};
-use crate::spec::{Schema, SchemaBuilder, SchemaRef};
+use crate::spec::{Schema, SchemaRef, TableMetadataBuilder};
 use crate::table::Table;
-use crate::transaction::snapshot::SnapshotProducer;
 use crate::transaction::{ActionCommit, TransactionAction};
 use crate::{TableRequirement, TableUpdate};
 
 /// UpdateSchemaAction is a transaction action for performing schema evolution to the table.
+///
+/// TODO(hjiang): Currently only drop column is supported, need to implement a few more schema evolution operation, i.e., add columns, rename columns, change nullability, update default, etc.
 pub struct UpdateSchemaAction {
     /// Schema update attributes.
     case_sensitive: bool,
@@ -90,13 +90,20 @@ impl UpdateSchemaAction {
             return Err(Error::new(
                 ErrorKind::PreconditionFailed,
                 format!(
-                    "Column '{}' is the table identifier, which canot be dropped.",
+                    "Column '{}' is the table identifier, which cannot be dropped.",
                     full_name
                 ),
             ));
         }
 
+        // Validate not all columns are dropped.
         self.deletes.insert(field.id);
+        if self.schema.field_id_to_fields().len() == self.deletes.len() {
+            return Err(Error::new(
+                ErrorKind::PreconditionFailed,
+                format!("Cannot delete all columns '{}' in the table.", full_name),
+            ));
+        }
 
         Ok(self)
     }
@@ -134,11 +141,108 @@ impl TransactionAction for UpdateSchemaAction {
         }
 
         let new_schema = self.get_updated_schema()?;
-        let new_schema_id = new_schema.schema_id();
         updates.push(TableUpdate::AddSchema { schema: new_schema });
         updates.push(TableUpdate::SetCurrentSchema {
-            schema_id: new_schema_id,
+            schema_id: TableMetadataBuilder::LAST_ADDED,
         });
         Ok(ActionCommit::new(updates, requirements))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::{HashMap, HashSet};
+    use std::sync::Arc;
+
+    use tempfile::TempDir;
+
+    use crate::io::FileIOBuilder;
+    use crate::spec::Schema;
+    use crate::transaction::tests::make_v2_table;
+    use crate::transaction::{ApplyTransactionAction, Transaction};
+    use crate::{Catalog, MemoryCatalog, NamespaceIdent, TableCreation, TableIdent};
+
+    /// Test util function to get [`TableCreation`].
+    async fn create_table(
+        catalog: &mut MemoryCatalog,
+        schema: Arc<Schema>,
+        warehouse_location: &str,
+    ) {
+        let table_name = "test1".to_string();
+        let namespace_ident = NamespaceIdent::from_vec(vec!["ns1".to_string()]).unwrap();
+        // let table_ident = TableIdent::new(namespace_ident.clone(), table_name.clone());
+        // let file_io = FileIOBuilder::new_fs_io().build().unwrap();
+
+        let table_creation = TableCreation::builder()
+            .name(table_name.clone())
+            .location(format!(
+                "{}/{}/{}",
+                warehouse_location,
+                namespace_ident.to_url_string(),
+                table_name
+            ))
+            .schema(schema.as_ref().clone())
+            .build();
+
+        catalog
+            .create_namespace(&namespace_ident, /*properties=*/ HashMap::new())
+            .await
+            .unwrap();
+        catalog
+            .create_table(&namespace_ident, table_creation)
+            .await
+            .unwrap();
+    }
+
+    #[test]
+    fn test_delete_empty_columns() {
+        let table = make_v2_table();
+        let tx = Transaction::new(&table);
+        let action = tx.update_schema();
+        assert!(action.deletes.is_empty());
+    }
+
+    #[test]
+    fn test_fail_to_delete_identifier_column() {
+        let table = make_v2_table();
+        let tx = Transaction::new(&table);
+        let mut action = tx.update_schema();
+        let res = action.delete_column(vec!["x".to_string()]);
+        assert!(res.is_err());
+    }
+
+    #[test]
+    fn test_delete_non_identifier_column() {
+        let table = make_v2_table();
+        let tx = Transaction::new(&table);
+        let mut action = tx.update_schema();
+        action.delete_column(vec!["z".to_string()]).unwrap();
+        assert_eq!(action.deletes, HashSet::from([(3)]));
+    }
+
+    // Test column deletion with memory catalog.
+    #[tokio::test]
+    async fn test_delete_columns_with_catalog() {
+        let file_io = FileIOBuilder::new_fs_io().build().unwrap();
+        let temp_dir = TempDir::new().unwrap();
+        let warehouse_location = temp_dir.path().to_str().unwrap().to_string();
+
+        let table = make_v2_table();
+        let mut memory_catalog = MemoryCatalog::new(file_io, Some(warehouse_location.clone()));
+        let schema = table.metadata().current_schema().clone();
+        create_table(&mut memory_catalog, schema, &warehouse_location).await;
+
+        let mut tx = Transaction::new(&table);
+        let mut action = tx.update_schema();
+        action.delete_column(vec!["z".to_string()]).unwrap();
+        tx = action.apply(tx).unwrap();
+
+        let table = tx.commit(&memory_catalog).await.unwrap();
+        let schema = table.metadata().current_schema();
+        assert!(schema.field_by_id(/*field_id=*/ 1).is_some());
+        assert!(schema.field_by_id(/*field_id=*/ 2).is_some());
+        assert!(schema.field_by_id(/*field_id=*/ 3).is_none());
+        assert_eq!(schema.highest_field_id(), 2);
+        assert_eq!(schema.identifier_field_ids().len(), 2);
+    }
+}