This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 528c4ab6ad Allow access to UDTF in `SessionContext` (#11071)
528c4ab6ad is described below
commit 528c4ab6ad5c974825ad623fc63eb300c203a356
Author: Heran Lin <[email protected]>
AuthorDate: Tue Jun 25 04:33:09 2024 +0800
Allow access to UDTF in `SessionContext` (#11071)
* get table function from session context
* support deregistering UDTF
* add `state_ref()` for `SessionContext`
---------
Co-authored-by: Shehab <[email protected]>
---
datafusion/core/src/datasource/function.rs | 5 ++++
datafusion/core/src/execution/context/mod.rs | 29 ++++++++++++++++++++--
datafusion/core/src/execution/session_state.rs | 14 +++++++++++
.../user_defined/user_defined_table_functions.rs | 15 +++++++++++
4 files changed, 61 insertions(+), 2 deletions(-)
diff --git a/datafusion/core/src/datasource/function.rs
b/datafusion/core/src/datasource/function.rs
index 2fd352ee4e..14bbc431f9 100644
--- a/datafusion/core/src/datasource/function.rs
+++ b/datafusion/core/src/datasource/function.rs
@@ -49,6 +49,11 @@ impl TableFunction {
&self.name
}
+ /// Get the implementation of the table function
+ pub fn function(&self) -> &Arc<dyn TableFunctionImpl> {
+ &self.fun
+ }
+
/// Get the function implementation and generate a table
pub fn create_table_provider(&self, args: &[Expr]) -> Result<Arc<dyn
TableProvider>> {
self.fun.call(args)
diff --git a/datafusion/core/src/execution/context/mod.rs
b/datafusion/core/src/execution/context/mod.rs
index 57e94cad06..52fc31de78 100644
--- a/datafusion/core/src/execution/context/mod.rs
+++ b/datafusion/core/src/execution/context/mod.rs
@@ -28,7 +28,7 @@ use crate::{
catalog::{CatalogProvider, CatalogProviderList, MemoryCatalogProvider},
dataframe::DataFrame,
datasource::{
- function::TableFunctionImpl,
+ function::{TableFunction, TableFunctionImpl},
listing::{ListingOptions, ListingTable, ListingTableConfig,
ListingTableUrl},
provider::TableProviderFactory,
},
@@ -52,7 +52,7 @@ use arrow::record_batch::RecordBatch;
use arrow_schema::Schema;
use datafusion_common::{
config::{ConfigExtension, TableOptions},
- exec_err, not_impl_err, plan_err,
+ exec_err, not_impl_err, plan_datafusion_err, plan_err,
tree_node::{TreeNodeRecursion, TreeNodeVisitor},
DFSchema, SchemaReference, TableReference,
};
@@ -928,6 +928,7 @@ impl SessionContext {
dropped |= self.state.write().deregister_udf(&stmt.name)?.is_some();
dropped |= self.state.write().deregister_udaf(&stmt.name)?.is_some();
dropped |= self.state.write().deregister_udwf(&stmt.name)?.is_some();
+ dropped |= self.state.write().deregister_udtf(&stmt.name)?.is_some();
// DROP FUNCTION IF EXISTS drops the specified function only if that
// function exists and in this way, it avoids error. While the DROP
FUNCTION
@@ -1008,6 +1009,11 @@ impl SessionContext {
self.state.write().deregister_udwf(name).ok();
}
+ /// Deregisters a UDTF within this context.
+ pub fn deregister_udtf(&self, name: &str) {
+ self.state.write().deregister_udtf(name).ok();
+ }
+
/// Creates a [`DataFrame`] for reading a data source.
///
/// For more control such as reading multiple files, you can use
@@ -1266,6 +1272,20 @@ impl SessionContext {
Ok(DataFrame::new(self.state(), plan))
}
+ /// Retrieves a [`TableFunction`] reference by name.
+ ///
+ /// Returns an error if no table function has been registered with the
provided name.
+ ///
+ /// [`register_udtf`]: SessionContext::register_udtf
+ pub fn table_function(&self, name: &str) -> Result<Arc<TableFunction>> {
+ self.state
+ .read()
+ .table_functions()
+ .get(name)
+ .cloned()
+ .ok_or_else(|| plan_datafusion_err!("Table function '{name}' not
found"))
+ }
+
/// Return a [`TableProvider`] for the specified table.
pub async fn table_provider<'a>(
&self,
@@ -1303,6 +1323,11 @@ impl SessionContext {
state
}
+ /// Get reference to [`SessionState`]
+ pub fn state_ref(&self) -> Arc<RwLock<SessionState>> {
+ self.state.clone()
+ }
+
/// Get weak reference to [`SessionState`]
pub fn state_weak_ref(&self) -> Weak<RwLock<SessionState>> {
Arc::downgrade(&self.state)
diff --git a/datafusion/core/src/execution/session_state.rs
b/datafusion/core/src/execution/session_state.rs
index 4173a8dcc4..60f67a5b23 100644
--- a/datafusion/core/src/execution/session_state.rs
+++ b/datafusion/core/src/execution/session_state.rs
@@ -834,6 +834,11 @@ impl SessionState {
&self.window_functions
}
+ /// Return reference to table_functions
+ pub fn table_functions(&self) -> &HashMap<String, Arc<TableFunction>> {
+ &self.table_functions
+ }
+
/// Return [SerializerRegistry] for extensions
pub fn serializer_registry(&self) -> Arc<dyn SerializerRegistry> {
self.serializer_registry.clone()
@@ -851,6 +856,15 @@ impl SessionState {
Arc::new(TableFunction::new(name.to_owned(), fun)),
);
}
+
+ /// Deregsiter a user defined table function
+ pub fn deregister_udtf(
+ &mut self,
+ name: &str,
+ ) -> datafusion_common::Result<Option<Arc<dyn TableFunctionImpl>>> {
+ let udtf = self.table_functions.remove(name);
+ Ok(udtf.map(|x| x.function().clone()))
+ }
}
struct SessionContextProvider<'a> {
diff --git a/datafusion/core/tests/user_defined/user_defined_table_functions.rs
b/datafusion/core/tests/user_defined/user_defined_table_functions.rs
index d3ddbed20d..1e8d30cab6 100644
--- a/datafusion/core/tests/user_defined/user_defined_table_functions.rs
+++ b/datafusion/core/tests/user_defined/user_defined_table_functions.rs
@@ -90,6 +90,21 @@ async fn test_simple_read_csv_udtf() -> Result<()> {
Ok(())
}
+#[tokio::test]
+async fn test_deregister_udtf() -> Result<()> {
+ let ctx = SessionContext::new();
+
+ ctx.register_udtf("read_csv", Arc::new(SimpleCsvTableFunc {}));
+
+ assert!(ctx.state().table_functions().contains_key("read_csv"));
+
+ ctx.deregister_udtf("read_csv");
+
+ assert!(!ctx.state().table_functions().contains_key("read_csv"));
+
+ Ok(())
+}
+
struct SimpleCsvTable {
schema: SchemaRef,
exprs: Vec<Expr>,
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
