diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index 4685f194fe29..04debf498aa9 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> { + self.state.read().expr_planners() + } + fn register_user_defined_sql_planner( &mut self, user_defined_sql_planner: Arc, diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs index ffaaa2df5e7e..ad557b12255c 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> { + self.user_defined_sql_planners.clone() + } + fn register_user_defined_sql_planner( &mut self, user_defined_sql_planner: Arc, 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 5e3c44c039ab..ae8a009c6292 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 c21ce3d21da1..24d61e6a8b72 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> { + vec![] + } } #[cfg(test)] diff --git a/datafusion/expr/src/registry.rs b/datafusion/expr/src/registry.rs index c276fe30f897..6a27c05bb451 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>; + /// 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) -> Result>> { Ok(self.udwfs.insert(udaf.name().into(), udaf)) } + + fn expr_planners(&self) -> Vec> { + vec![] + } } diff --git a/datafusion/proto/src/bytes/mod.rs b/datafusion/proto/src/bytes/mod.rs index 901aa2455e16..83210cb4e41f 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> { + 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 4bf2bb3d7b79..075993e2ba76 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) -> Result>> { plan_err!("No function registry provided to deserialize, so can not deserialize User Defined Window Function '{}'", udwf.inner().name()) } + + fn expr_planners(&self) -> Vec> { + vec![] + } }