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 9f8ba6ab68 Add user_defined_sql_planners(..) to FunctionRegistry
(#11296)
9f8ba6ab68 is described below
commit 9f8ba6ab686434d6abfae42d2b2b2819e4bbbd9d
Author: Bruce Ritchie <[email protected]>
AuthorDate: Sun Jul 7 13:03:10 2024 -0400
Add user_defined_sql_planners(..) to FunctionRegistry (#11296)
* Add user_defined_sql_planners(..) to FunctionRegistry
* Adding simple test for user_defined_sql_planners
* Renamed user_defined_sql_planners to expr_planners
---
datafusion/core/src/execution/context/mod.rs | 4 ++++
datafusion/core/src/execution/session_state.rs | 4 ++++
.../core/tests/user_defined/user_defined_scalar_functions.rs | 11 +++++++++++
datafusion/execution/src/task.rs | 10 +++++++---
datafusion/expr/src/registry.rs | 7 +++++++
datafusion/proto/src/bytes/mod.rs | 5 +++++
datafusion/proto/src/bytes/registry.rs | 5 +++++
7 files changed, 43 insertions(+), 3 deletions(-)
diff --git a/datafusion/core/src/execution/context/mod.rs
b/datafusion/core/src/execution/context/mod.rs
index 4685f194fe..04debf498a 100644
--- a/datafusion/core/src/execution/context/mod.rs
+++ b/datafusion/core/src/execution/context/mod.rs
@@ -1392,6 +1392,10 @@ impl FunctionRegistry for SessionContext {
self.state.write().register_function_rewrite(rewrite)
}
+ fn expr_planners(&self) -> Vec<Arc<dyn UserDefinedSQLPlanner>> {
+ self.state.read().expr_planners()
+ }
+
fn register_user_defined_sql_planner(
&mut self,
user_defined_sql_planner: Arc<dyn UserDefinedSQLPlanner>,
diff --git a/datafusion/core/src/execution/session_state.rs
b/datafusion/core/src/execution/session_state.rs
index ffaaa2df5e..ad557b1225 100644
--- a/datafusion/core/src/execution/session_state.rs
+++ b/datafusion/core/src/execution/session_state.rs
@@ -1183,6 +1183,10 @@ impl FunctionRegistry for SessionState {
Ok(())
}
+ fn expr_planners(&self) -> Vec<Arc<dyn UserDefinedSQLPlanner>> {
+ self.user_defined_sql_planners.clone()
+ }
+
fn register_user_defined_sql_planner(
&mut self,
user_defined_sql_planner: Arc<dyn UserDefinedSQLPlanner>,
diff --git
a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs
b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs
index 5e3c44c039..ae8a009c62 100644
--- a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs
+++ b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs
@@ -571,6 +571,17 @@ async fn test_user_defined_functions_cast_to_i64() ->
Result<()> {
Ok(())
}
+#[tokio::test]
+async fn test_user_defined_sql_functions() -> Result<()> {
+ let ctx = SessionContext::new();
+
+ let sql_planners = ctx.expr_planners();
+
+ assert!(!sql_planners.is_empty());
+
+ Ok(())
+}
+
#[tokio::test]
async fn deregister_udf() -> Result<()> {
let cast2i64 = ScalarUDF::from(CastToI64UDF::new());
diff --git a/datafusion/execution/src/task.rs b/datafusion/execution/src/task.rs
index c21ce3d21d..24d61e6a8b 100644
--- a/datafusion/execution/src/task.rs
+++ b/datafusion/execution/src/task.rs
@@ -20,15 +20,15 @@ use std::{
sync::Arc,
};
-use datafusion_common::{plan_datafusion_err, DataFusionError, Result};
-use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF};
-
use crate::{
config::SessionConfig,
memory_pool::MemoryPool,
registry::FunctionRegistry,
runtime_env::{RuntimeConfig, RuntimeEnv},
};
+use datafusion_common::{plan_datafusion_err, DataFusionError, Result};
+use datafusion_expr::planner::UserDefinedSQLPlanner;
+use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF};
/// Task Execution Context
///
@@ -191,6 +191,10 @@ impl FunctionRegistry for TaskContext {
});
Ok(self.scalar_functions.insert(udf.name().into(), udf))
}
+
+ fn expr_planners(&self) -> Vec<Arc<dyn UserDefinedSQLPlanner>> {
+ vec![]
+ }
}
#[cfg(test)]
diff --git a/datafusion/expr/src/registry.rs b/datafusion/expr/src/registry.rs
index c276fe30f8..6a27c05bb4 100644
--- a/datafusion/expr/src/registry.rs
+++ b/datafusion/expr/src/registry.rs
@@ -110,6 +110,9 @@ pub trait FunctionRegistry {
not_impl_err!("Registering FunctionRewrite")
}
+ /// Set of all registered [`UserDefinedSQLPlanner`]s
+ fn expr_planners(&self) -> Vec<Arc<dyn UserDefinedSQLPlanner>>;
+
/// Registers a new [`UserDefinedSQLPlanner`] with the registry.
fn register_user_defined_sql_planner(
&mut self,
@@ -192,4 +195,8 @@ impl FunctionRegistry for MemoryFunctionRegistry {
fn register_udwf(&mut self, udaf: Arc<WindowUDF>) ->
Result<Option<Arc<WindowUDF>>> {
Ok(self.udwfs.insert(udaf.name().into(), udaf))
}
+
+ fn expr_planners(&self) -> Vec<Arc<dyn UserDefinedSQLPlanner>> {
+ vec![]
+ }
}
diff --git a/datafusion/proto/src/bytes/mod.rs
b/datafusion/proto/src/bytes/mod.rs
index 901aa2455e..83210cb4e4 100644
--- a/datafusion/proto/src/bytes/mod.rs
+++ b/datafusion/proto/src/bytes/mod.rs
@@ -39,6 +39,7 @@ use std::sync::Arc;
use datafusion::execution::registry::FunctionRegistry;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::prelude::SessionContext;
+use datafusion_expr::planner::UserDefinedSQLPlanner;
mod registry;
@@ -165,6 +166,10 @@ impl Serializeable for Expr {
"register_udwf called in Placeholder Registry!"
)
}
+
+ fn expr_planners(&self) -> Vec<Arc<dyn UserDefinedSQLPlanner>> {
+ vec![]
+ }
}
Expr::from_bytes_with_registry(&bytes, &PlaceHolderRegistry)?;
diff --git a/datafusion/proto/src/bytes/registry.rs
b/datafusion/proto/src/bytes/registry.rs
index 4bf2bb3d7b..075993e2ba 100644
--- a/datafusion/proto/src/bytes/registry.rs
+++ b/datafusion/proto/src/bytes/registry.rs
@@ -20,6 +20,7 @@ use std::{collections::HashSet, sync::Arc};
use datafusion::execution::registry::FunctionRegistry;
use datafusion_common::plan_err;
use datafusion_common::Result;
+use datafusion_expr::planner::UserDefinedSQLPlanner;
use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF};
/// A default [`FunctionRegistry`] registry that does not resolve any
@@ -54,4 +55,8 @@ impl FunctionRegistry for NoRegistry {
fn register_udwf(&mut self, udwf: Arc<WindowUDF>) ->
Result<Option<Arc<WindowUDF>>> {
plan_err!("No function registry provided to deserialize, so can not
deserialize User Defined Window Function '{}'", udwf.inner().name())
}
+
+ fn expr_planners(&self) -> Vec<Arc<dyn UserDefinedSQLPlanner>> {
+ vec![]
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
