From 7ba767f0e9c287b99c1decccc10911f3ec222216 Mon Sep 17 00:00:00 2001 From: Shawn Chang Date: Thu, 15 Jan 2026 13:09:39 -0800 Subject: [PATCH] Add support for DROP TABLE --- crates/integrations/datafusion/src/schema.rs | 76 +++++++++++++++++++ .../testdata/schedules/df_test.toml | 4 + .../testdata/slts/df_test/drop_table.slt | 63 +++++++++++++++ 3 files changed, 143 insertions(+) create mode 100644 crates/sqllogictest/testdata/slts/df_test/drop_table.slt diff --git a/crates/integrations/datafusion/src/schema.rs b/crates/integrations/datafusion/src/schema.rs index 022964ba6d..ab8f89a5ad 100644 --- a/crates/integrations/datafusion/src/schema.rs +++ b/crates/integrations/datafusion/src/schema.rs @@ -33,6 +33,7 @@ use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, Result, TableCreation}; use crate::table::IcebergTableProvider; use crate::to_datafusion_error; +use iceberg::TableIdent; /// Represents a [`SchemaProvider`] for the Iceberg [`Catalog`], managing /// access to table providers within a specific namespace. @@ -211,6 +212,43 @@ impl SchemaProvider for IcebergSchemaProvider { DataFusionError::Execution(format!("Failed to create Iceberg table: {e}")) })? } + + fn deregister_table(&self, name: &str) -> DFResult>> { + // Check if table exists + if !self.table_exist(name) { + return Ok(None); + } + + let catalog = self.catalog.clone(); + let namespace = self.namespace.clone(); + let tables = self.tables.clone(); + let table_name = name.to_string(); + + // Use tokio's spawn_blocking to handle the async work on a blocking thread pool + let result = tokio::task::spawn_blocking(move || { + let rt = tokio::runtime::Handle::current(); + rt.block_on(async move { + let table_ident = TableIdent::new(namespace, table_name.clone()); + + // Drop the table from the Iceberg catalog + catalog + .drop_table(&table_ident) + .await + .map_err(to_datafusion_error)?; + + // Remove from local cache and return the removed provider + let removed = tables + .remove(&table_name) + .map(|(_, table)| table as Arc); + + Ok(removed) + }) + }); + + futures::executor::block_on(result).map_err(|e| { + DataFusionError::Execution(format!("Failed to drop Iceberg table: {e}")) + })? + } } /// Verifies that a table provider contains no data by scanning with LIMIT 1. @@ -368,4 +406,42 @@ mod tests { "Expected error about table already existing, got: {err}", ); } + + #[tokio::test] + async fn test_deregister_table_succeeds() { + let (schema_provider, _temp_dir) = create_test_schema_provider().await; + + // Create and register an empty table + let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new( + "id", + DataType::Int32, + false, + )])); + + let empty_batch = RecordBatch::new_empty(arrow_schema.clone()); + let mem_table = MemTable::try_new(arrow_schema, vec![vec![empty_batch]]).unwrap(); + + // Register the table + let result = schema_provider.register_table("drop_me".to_string(), Arc::new(mem_table)); + assert!(result.is_ok()); + assert!(schema_provider.table_exist("drop_me")); + + // Deregister the table + let result = schema_provider.deregister_table("drop_me"); + assert!(result.is_ok()); + assert!(result.unwrap().is_some()); + + // Verify the table no longer exists + assert!(!schema_provider.table_exist("drop_me")); + } + + #[tokio::test] + async fn test_deregister_nonexistent_table_returns_none() { + let (schema_provider, _temp_dir) = create_test_schema_provider().await; + + // Attempt to deregister a table that doesn't exist + let result = schema_provider.deregister_table("nonexistent"); + assert!(result.is_ok()); + assert!(result.unwrap().is_none()); + } } diff --git a/crates/sqllogictest/testdata/schedules/df_test.toml b/crates/sqllogictest/testdata/schedules/df_test.toml index 1d7f42c8d4..5fc6bfa945 100644 --- a/crates/sqllogictest/testdata/schedules/df_test.toml +++ b/crates/sqllogictest/testdata/schedules/df_test.toml @@ -29,3 +29,7 @@ slt = "df_test/create_table.slt" [[steps]] engine = "df" slt = "df_test/insert_into.slt" + +[[steps]] +engine = "df" +slt = "df_test/drop_table.slt" diff --git a/crates/sqllogictest/testdata/slts/df_test/drop_table.slt b/crates/sqllogictest/testdata/slts/df_test/drop_table.slt new file mode 100644 index 0000000000..c721c84e11 --- /dev/null +++ b/crates/sqllogictest/testdata/slts/df_test/drop_table.slt @@ -0,0 +1,63 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Test DROP TABLE functionality + +# Create a table to drop +statement ok +CREATE TABLE default.default.table_to_drop (id INT NOT NULL, name STRING) + +# Verify the table exists +query IT rowsort +SELECT * FROM default.default.table_to_drop +---- + +# Insert some data +query I +INSERT INTO default.default.table_to_drop VALUES (1, 'Test') +---- +1 + +# Verify data exists +query IT rowsort +SELECT * FROM default.default.table_to_drop +---- +1 Test + +# Drop the table +statement ok +DROP TABLE default.default.table_to_drop + +# Verify the table no longer exists (should error) +statement error +SELECT * FROM default.default.table_to_drop + +# Test DROP TABLE IF EXISTS on non-existent table (should not error) +statement ok +DROP TABLE IF EXISTS default.default.nonexistent_table + +# Create another table for IF EXISTS test +statement ok +CREATE TABLE default.default.another_table (id INT NOT NULL) + +# Drop with IF EXISTS +statement ok +DROP TABLE IF EXISTS default.default.another_table + +# Verify it's gone +statement error +SELECT * FROM default.default.another_table