diff --git a/rust/datafusion/src/execution/context.rs b/rust/datafusion/src/execution/context.rs index 4b8a4f71fe6..766d74ed58e 100644 --- a/rust/datafusion/src/execution/context.rs +++ b/rust/datafusion/src/execution/context.rs @@ -287,6 +287,18 @@ impl ExecutionContext { .insert(name.to_string(), provider.into()); } + /// Deregisters the named table. + /// + /// Returns true if the table was successfully de-reregistered. + pub fn deregister_table(&mut self, name: &str) -> bool { + self.state + .lock() + .unwrap() + .datasources + .remove(&name.to_string()) + .is_some() + } + /// Retrieves a DataFrame representing a table previously registered by calling the /// register_table function. /// @@ -723,6 +735,21 @@ mod tests { Ok(()) } + #[tokio::test] + async fn register_deregister() -> Result<()> { + let tmp_dir = TempDir::new()?; + let partition_count = 4; + let mut ctx = create_ctx(&tmp_dir, partition_count)?; + + let provider = test::create_table_dual(); + ctx.register_table("dual", provider); + + assert_eq!(ctx.deregister_table("dual"), true); + assert_eq!(ctx.deregister_table("dual"), false); + + Ok(()) + } + #[tokio::test] async fn parallel_query_with_filter() -> Result<()> { let tmp_dir = TempDir::new()?; @@ -1668,6 +1695,8 @@ mod tests { assert_eq!(a.value(i) + b.value(i), sum.value(i)); } + ctx.deregister_table("t"); + Ok(()) }