feat: pass SessionState not SessionConfig to FunctionFactory::create (a…

…pache#9837)
viirya · Mar 28, 2024 · 6f9948b · 6f9948b
1 parent b6915f5
commit 6f9948b
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 6 deletions.
diff --git a/datafusion-examples/examples/function_factory.rs b/datafusion-examples/examples/function_factory.rs
@@ -16,8 +16,9 @@
 // under the License.
 
 use datafusion::error::Result;
-use datafusion::execution::config::SessionConfig;
-use datafusion::execution::context::{FunctionFactory, RegisterFunction, SessionContext};
+use datafusion::execution::context::{
+    FunctionFactory, RegisterFunction, SessionContext, SessionState,
+};
 use datafusion_common::tree_node::{Transformed, TreeNode};
 use datafusion_common::{exec_err, internal_err, DataFusionError};
 use datafusion_expr::simplify::ExprSimplifyResult;
@@ -91,7 +92,7 @@ impl FunctionFactory for CustomFunctionFactory {
     /// the function instance.
     async fn create(
         &self,
-        _state: &SessionConfig,
+        _state: &SessionState,
         statement: CreateFunction,
     ) -> Result<RegisterFunction> {
         let f: ScalarFunctionWrapper = statement.try_into()?;

diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs
@@ -794,7 +794,7 @@ impl SessionContext {
             let function_factory = &state.function_factory;
 
             match function_factory {
-                Some(f) => f.create(state.config(), stmt).await?,
+                Some(f) => f.create(&state, stmt).await?,
                 _ => Err(DataFusionError::Configuration(
                     "Function factory has not been configured".into(),
                 ))?,
@@ -1288,7 +1288,7 @@ pub trait FunctionFactory: Sync + Send {
     /// Handles creation of user defined function specified in [CreateFunction] statement
     async fn create(
         &self,
-        state: &SessionConfig,
+        state: &SessionState,
         statement: CreateFunction,
     ) -> Result<RegisterFunction>;
 }

diff --git a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs
@@ -747,7 +747,7 @@ struct CustomFunctionFactory {}
 impl FunctionFactory for CustomFunctionFactory {
     async fn create(
         &self,
-        _state: &SessionConfig,
+        _state: &SessionState,
         statement: CreateFunction,
     ) -> Result<RegisterFunction> {
         let f: ScalarFunctionWrapper = statement.try_into()?;