Skip to content

Commit 0b93b6a

Browse files
committed
remove's gone from the trait ...
... `DROP FUNCTION` will look for function name in all available registries (udf, udaf, udwf). `remove` may be necessary if UDaF and UDwF do not get `simplify` method from #9304.
1 parent 45c7edb commit 0b93b6a

File tree

2 files changed

+41
-62
lines changed

2 files changed

+41
-62
lines changed

datafusion/core/src/execution/context/mod.rs

Lines changed: 26 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -827,20 +827,23 @@ impl SessionContext {
827827
}
828828

829829
async fn drop_function(&self, stmt: DropFunction) -> Result<DataFrame> {
830-
let _function = {
831-
let state = self.state.read().clone();
832-
let function_factory = &state.function_factory;
833-
834-
match function_factory {
835-
Some(f) => f.remove(state.config(), stmt).await?,
836-
None => Err(DataFusionError::Configuration(
837-
"Function factory has not been configured".into(),
838-
))?,
839-
}
840-
};
841-
842-
// TODO: Once we have unregister UDF we need to implement it here
843-
self.return_empty_dataframe()
830+
// we don't know function type at this point
831+
// decision has been made to drop all functions
832+
let mut dropped = false;
833+
dropped |= self.state.write().deregister_udf(&stmt.name)?.is_some();
834+
dropped |= self.state.write().deregister_udaf(&stmt.name)?.is_some();
835+
dropped |= self.state.write().deregister_udwf(&stmt.name)?.is_some();
836+
837+
// DROP FUNCTION IF EXISTS drops the specified function only if that
838+
// function exists and in this way, it avoids error. While the DROP FUNCTION
839+
// statement also performs the same function, it throws an
840+
// error if the function does not exist.
841+
842+
if !stmt.if_exists && !dropped {
843+
Err(DataFusionError::Execution("Function does not exist".into()))
844+
} else {
845+
self.return_empty_dataframe()
846+
}
844847
}
845848

846849
/// Registers a variable provider within this context.
@@ -1310,18 +1313,9 @@ impl QueryPlanner for DefaultQueryPlanner {
13101313
.await
13111314
}
13121315
}
1313-
/// Crates and registers a function from [CreateFunction] statement
1314-
///
1315-
/// It is intended to handle `CREATE FUNCTION` statements
1316-
/// and interact with [SessionState] to registers new udfs.
1317-
///
1318-
/// Datafusion `SQL` dialect does not support `CREATE FUNCTION`
1319-
/// in generic dialect, so dialect should be changed to `PostgreSQL`
1320-
///
1321-
/// ```rust, no_run
1322-
/// # use datafusion::execution::config::SessionConfig;
1323-
/// SessionConfig::new().set_str("datafusion.sql_parser.dialect", "PostgreSQL");
1324-
/// ```
1316+
/// A pluggable interface to handle `CREATE FUNCTION` statements
1317+
/// and interact with [SessionState] to registers new udf, udaf or udwf.
1318+
13251319
#[async_trait]
13261320
pub trait FunctionFactory: Sync + Send {
13271321
/// Handles creation of user defined function specified in [CreateFunction] statement
@@ -1331,13 +1325,13 @@ pub trait FunctionFactory: Sync + Send {
13311325
statement: CreateFunction,
13321326
) -> Result<RegisterFunction>;
13331327

1334-
/// Drops user defined function from [SessionState]
1328+
// Drops user defined function from [SessionState]
13351329
// Naming it `drop` would make more sense but its already occupied in rust
1336-
async fn remove(
1337-
&self,
1338-
state: &SessionConfig,
1339-
statement: DropFunction,
1340-
) -> Result<RegisterFunction>;
1330+
// async fn remove(
1331+
// &self,
1332+
// state: &SessionConfig,
1333+
// statement: DropFunction,
1334+
// ) -> Result<RegisterFunction>;
13411335
}
13421336

13431337
/// Type of function to create

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

Lines changed: 15 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,8 @@ use datafusion_common::{
3131
};
3232
use datafusion_execution::runtime_env::{RuntimeConfig, RuntimeEnv};
3333
use datafusion_expr::{
34-
create_udaf, create_udf, Accumulator, ColumnarValue, CreateFunction, DropFunction,
35-
ExprSchemable, LogicalPlanBuilder, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
34+
create_udaf, create_udf, Accumulator, ColumnarValue, CreateFunction, ExprSchemable,
35+
LogicalPlanBuilder, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
3636
};
3737
use parking_lot::Mutex;
3838
use rand::{thread_rng, Rng};
@@ -693,28 +693,10 @@ impl FunctionFactory for MockFunctionFactory {
693693

694694
Ok(RegisterFunction::Scalar(Arc::new(mock_udf)))
695695
}
696-
697-
async fn remove(
698-
&self,
699-
_config: &SessionConfig,
700-
_statement: DropFunction,
701-
) -> datafusion::error::Result<RegisterFunction> {
702-
// TODO: I don't like that remove returns RegisterFunction
703-
// we have to keep two states in FunctionFactory iml and
704-
// SessionState
705-
//
706-
// It would be better to return (function_name, function type) tuple
707-
//
708-
// at the moment state does not support unregister user defined functions
709-
710-
Err(datafusion_common::DataFusionError::NotImplemented(
711-
"remove function has not been implemented".into(),
712-
))
713-
}
714696
}
715697

716698
#[tokio::test]
717-
async fn create_scalar_function_from_sql_statement() {
699+
async fn create_scalar_function_from_sql_statement() -> Result<()> {
718700
let function_factory = Arc::new(MockFunctionFactory::default());
719701
let runtime_config = RuntimeConfig::new();
720702
let runtime_environment = RuntimeEnv::new(runtime_config).unwrap();
@@ -732,23 +714,26 @@ async fn create_scalar_function_from_sql_statement() {
732714
RETURNS DOUBLE
733715
RETURN $1 + $2
734716
"#;
735-
let _ = ctx.sql(sql).await.unwrap();
717+
let _ = ctx.sql(sql).await?;
736718

737-
ctx.sql("select better_add(2.0, 2.0)")
738-
.await
739-
.unwrap()
740-
.show()
741-
.await
742-
.unwrap();
719+
ctx.sql("select better_add(2.0, 2.0)").await?.show().await?;
743720

744721
// check if we sql expr has been converted to datafusion expr
745722
let captured_expression = function_factory.captured_expr.lock().clone().unwrap();
746723

747724
// is there some better way to test this
748725
assert_eq!("$1 + $2", captured_expression.to_string());
749726

750-
// no support at the moment
751-
// ctx.sql("drop function better_add").await.unwrap();
727+
// statement drops function
728+
assert!(ctx.sql("drop function better_add").await.is_ok());
729+
// no function, it panics
730+
assert!(ctx.sql("drop function better_add").await.is_err());
731+
// no function, it dies not care
732+
assert!(ctx.sql("drop function if exists better_add").await.is_ok());
733+
// query should fail as there is no function
734+
assert!(ctx.sql("select better_add(2.0, 2.0)").await.is_err());
735+
736+
Ok(())
752737
}
753738

754739
fn create_udf_context() -> SessionContext {

0 commit comments

Comments
 (0)