Skip to content

Commit 05651e2

Browse files
authored
feat(datafusion): Add support for DROP TABLE (#2033)
## Which issue does this PR close? - Closes #2060 ## What changes are included in this PR? - Support DROP TABLE syntax by implementing `SchemaProvider::deregister_table` ## Are these changes tested? Added sqllogictests
1 parent 20ce7a5 commit 05651e2

3 files changed

Lines changed: 142 additions & 1 deletion

File tree

crates/integrations/datafusion/src/schema.rs

Lines changed: 75 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ use futures::StreamExt;
2929
use futures::future::try_join_all;
3030
use iceberg::arrow::arrow_schema_to_schema_auto_assign_ids;
3131
use iceberg::inspect::MetadataTableType;
32-
use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, Result, TableCreation};
32+
use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, Result, TableCreation, TableIdent};
3333

3434
use crate::table::IcebergTableProvider;
3535
use crate::to_datafusion_error;
@@ -211,6 +211,42 @@ impl SchemaProvider for IcebergSchemaProvider {
211211
DataFusionError::Execution(format!("Failed to create Iceberg table: {e}"))
212212
})?
213213
}
214+
215+
fn deregister_table(&self, name: &str) -> DFResult<Option<Arc<dyn TableProvider>>> {
216+
// Check if table exists
217+
if !self.table_exist(name) {
218+
return Ok(None);
219+
}
220+
221+
let catalog = self.catalog.clone();
222+
let namespace = self.namespace.clone();
223+
let tables = self.tables.clone();
224+
let table_name = name.to_string();
225+
226+
// Use tokio's spawn_blocking to handle the async work on a blocking thread pool
227+
let result = tokio::task::spawn_blocking(move || {
228+
let rt = tokio::runtime::Handle::current();
229+
rt.block_on(async move {
230+
let table_ident = TableIdent::new(namespace, table_name.clone());
231+
232+
// Drop the table from the Iceberg catalog
233+
catalog
234+
.drop_table(&table_ident)
235+
.await
236+
.map_err(to_datafusion_error)?;
237+
238+
// Remove from local cache and return the removed provider
239+
let removed = tables
240+
.remove(&table_name)
241+
.map(|(_, table)| table as Arc<dyn TableProvider>);
242+
243+
Ok(removed)
244+
})
245+
});
246+
247+
futures::executor::block_on(result)
248+
.map_err(|e| DataFusionError::Execution(format!("Failed to drop Iceberg table: {e}")))?
249+
}
214250
}
215251

216252
/// Verifies that a table provider contains no data by scanning with LIMIT 1.
@@ -368,4 +404,42 @@ mod tests {
368404
"Expected error about table already existing, got: {err}",
369405
);
370406
}
407+
408+
#[tokio::test]
409+
async fn test_deregister_table_succeeds() {
410+
let (schema_provider, _temp_dir) = create_test_schema_provider().await;
411+
412+
// Create and register an empty table
413+
let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
414+
"id",
415+
DataType::Int32,
416+
false,
417+
)]));
418+
419+
let empty_batch = RecordBatch::new_empty(arrow_schema.clone());
420+
let mem_table = MemTable::try_new(arrow_schema, vec![vec![empty_batch]]).unwrap();
421+
422+
// Register the table
423+
let result = schema_provider.register_table("drop_me".to_string(), Arc::new(mem_table));
424+
assert!(result.is_ok());
425+
assert!(schema_provider.table_exist("drop_me"));
426+
427+
// Deregister the table
428+
let result = schema_provider.deregister_table("drop_me");
429+
assert!(result.is_ok());
430+
assert!(result.unwrap().is_some());
431+
432+
// Verify the table no longer exists
433+
assert!(!schema_provider.table_exist("drop_me"));
434+
}
435+
436+
#[tokio::test]
437+
async fn test_deregister_nonexistent_table_returns_none() {
438+
let (schema_provider, _temp_dir) = create_test_schema_provider().await;
439+
440+
// Attempt to deregister a table that doesn't exist
441+
let result = schema_provider.deregister_table("nonexistent");
442+
assert!(result.is_ok());
443+
assert!(result.unwrap().is_none());
444+
}
371445
}

crates/sqllogictest/testdata/schedules/df_test.toml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,3 +33,7 @@ slt = "df_test/insert_into.slt"
3333
[[steps]]
3434
engine = "df"
3535
slt = "df_test/binary_predicate_pushdown.slt"
36+
37+
[[steps]]
38+
engine = "df"
39+
slt = "df_test/drop_table.slt"
Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
# Test DROP TABLE functionality
19+
20+
# Create a table to drop
21+
statement ok
22+
CREATE TABLE default.default.table_to_drop (id INT NOT NULL, name STRING)
23+
24+
# Verify the table exists
25+
query IT rowsort
26+
SELECT * FROM default.default.table_to_drop
27+
----
28+
29+
# Insert some data
30+
query I
31+
INSERT INTO default.default.table_to_drop VALUES (1, 'Test')
32+
----
33+
1
34+
35+
# Verify data exists
36+
query IT rowsort
37+
SELECT * FROM default.default.table_to_drop
38+
----
39+
1 Test
40+
41+
# Drop the table
42+
statement ok
43+
DROP TABLE default.default.table_to_drop
44+
45+
# Verify the table no longer exists (should error)
46+
statement error
47+
SELECT * FROM default.default.table_to_drop
48+
49+
# Test DROP TABLE IF EXISTS on non-existent table (should not error)
50+
statement ok
51+
DROP TABLE IF EXISTS default.default.nonexistent_table
52+
53+
# Create another table for IF EXISTS test
54+
statement ok
55+
CREATE TABLE default.default.another_table (id INT NOT NULL)
56+
57+
# Drop with IF EXISTS
58+
statement ok
59+
DROP TABLE IF EXISTS default.default.another_table
60+
61+
# Verify it's gone
62+
statement error
63+
SELECT * FROM default.default.another_table

0 commit comments

Comments
 (0)