diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index 2841f373ab6d..b6ea47ecacf9 100644 --- a/datafusion/core/src/execution/context/mod.rs +++ b/datafusion/core/src/execution/context/mod.rs @@ -827,20 +827,23 @@ impl SessionContext { } async fn drop_function(&self, stmt: DropFunction) -> Result { - let _function = { - let state = self.state.read().clone(); - let function_factory = &state.function_factory; - - match function_factory { - Some(f) => f.remove(state.config(), stmt).await?, - None => Err(DataFusionError::Configuration( - "Function factory has not been configured".into(), - ))?, - } - }; - - // TODO: Once we have unregister UDF we need to implement it here - self.return_empty_dataframe() + // we don't know function type at this point + // decision has been made to drop all functions + let mut dropped = false; + 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(); + + // DROP FUNCTION IF EXISTS drops the specified function only if that + // function exists and in this way, it avoids error. While the DROP FUNCTION + // statement also performs the same function, it throws an + // error if the function does not exist. + + if !stmt.if_exists && !dropped { + Err(DataFusionError::Execution("Function does not exist".into())) + } else { + self.return_empty_dataframe() + } } /// Registers a variable provider within this context. @@ -1310,18 +1313,9 @@ impl QueryPlanner for DefaultQueryPlanner { .await } } -/// Crates and registers a function from [CreateFunction] statement -/// -/// It is intended to handle `CREATE FUNCTION` statements -/// and interact with [SessionState] to registers new udfs. -/// -/// Datafusion `SQL` dialect does not support `CREATE FUNCTION` -/// in generic dialect, so dialect should be changed to `PostgreSQL` -/// -/// ```rust, no_run -/// # use datafusion::execution::config::SessionConfig; -/// SessionConfig::new().set_str("datafusion.sql_parser.dialect", "PostgreSQL"); -/// ``` +/// A pluggable interface to handle `CREATE FUNCTION` statements +/// and interact with [SessionState] to registers new udf, udaf or udwf. + #[async_trait] pub trait FunctionFactory: Sync + Send { /// Handles creation of user defined function specified in [CreateFunction] statement @@ -1330,14 +1324,6 @@ pub trait FunctionFactory: Sync + Send { state: &SessionConfig, statement: CreateFunction, ) -> Result; - - /// Drops user defined function from [SessionState] - // Naming it `drop` would make more sense but its already occupied in rust - async fn remove( - &self, - state: &SessionConfig, - statement: DropFunction, - ) -> Result; } /// Type of function to create 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 d178fbea4c00..f6bea0160f1f 100644 --- a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs @@ -31,8 +31,8 @@ use datafusion_common::{ }; use datafusion_execution::runtime_env::{RuntimeConfig, RuntimeEnv}; use datafusion_expr::{ - create_udaf, create_udf, Accumulator, ColumnarValue, CreateFunction, DropFunction, - ExprSchemable, LogicalPlanBuilder, ScalarUDF, ScalarUDFImpl, Signature, Volatility, + create_udaf, create_udf, Accumulator, ColumnarValue, CreateFunction, ExprSchemable, + LogicalPlanBuilder, ScalarUDF, ScalarUDFImpl, Signature, Volatility, }; use parking_lot::Mutex; use rand::{thread_rng, Rng}; @@ -693,28 +693,10 @@ impl FunctionFactory for MockFunctionFactory { Ok(RegisterFunction::Scalar(Arc::new(mock_udf))) } - - async fn remove( - &self, - _config: &SessionConfig, - _statement: DropFunction, - ) -> datafusion::error::Result { - // TODO: I don't like that remove returns RegisterFunction - // we have to keep two states in FunctionFactory iml and - // SessionState - // - // It would be better to return (function_name, function type) tuple - // - // at the moment state does not support unregister user defined functions - - Err(datafusion_common::DataFusionError::NotImplemented( - "remove function has not been implemented".into(), - )) - } } #[tokio::test] -async fn create_scalar_function_from_sql_statement() { +async fn create_scalar_function_from_sql_statement() -> Result<()> { let function_factory = Arc::new(MockFunctionFactory::default()); let runtime_config = RuntimeConfig::new(); let runtime_environment = RuntimeEnv::new(runtime_config).unwrap(); @@ -732,14 +714,9 @@ async fn create_scalar_function_from_sql_statement() { RETURNS DOUBLE RETURN $1 + $2 "#; - let _ = ctx.sql(sql).await.unwrap(); + let _ = ctx.sql(sql).await?; - ctx.sql("select better_add(2.0, 2.0)") - .await - .unwrap() - .show() - .await - .unwrap(); + ctx.sql("select better_add(2.0, 2.0)").await?.show().await?; // check if we sql expr has been converted to datafusion expr let captured_expression = function_factory.captured_expr.lock().clone().unwrap(); @@ -747,8 +724,16 @@ async fn create_scalar_function_from_sql_statement() { // is there some better way to test this assert_eq!("$1 + $2", captured_expression.to_string()); - // no support at the moment - // ctx.sql("drop function better_add").await.unwrap(); + // statement drops function + assert!(ctx.sql("drop function better_add").await.is_ok()); + // no function, it panics + assert!(ctx.sql("drop function better_add").await.is_err()); + // no function, it dies not care + assert!(ctx.sql("drop function if exists better_add").await.is_ok()); + // query should fail as there is no function + assert!(ctx.sql("select better_add(2.0, 2.0)").await.is_err()); + + Ok(()) } fn create_udf_context() -> SessionContext {