You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "milenkovicm (via GitHub)" <gi...@apache.org> on 2024/02/24 13:40:00 UTC
[PR] Add plugable function factory [arrow-datafusion]
milenkovicm opened a new pull request, #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333
## Which issue does this PR close?
Closes #9332.
## Rationale for this change
explained in #9332
## What changes are included in this PR?
this is minimum viable product, with some open question which i hope to resolve in the process of PR review (note TODOs)
## Are these changes tested?
There is a test in place but would add more if community aggress to accept the change
## Are there any user-facing changes?
There are using facing changes, but they are backward compatible. Would update docs if community aggress to accept the change
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1509952063
##########
datafusion/expr/src/logical_plan/ddl.rs:
##########
@@ -303,3 +321,66 @@ pub struct DropCatalogSchema {
/// Dummy schema
pub schema: DFSchemaRef,
}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunction {
+ // TODO: There is open question should we expose sqlparser types or redefine them here?
+ // At the moment it make more sense to expose sqlparser types and leave
+ // user to convert them as needed
+ pub or_replace: bool,
+ pub temporary: bool,
+ pub name: String,
+ pub args: Option<Vec<OperateFunctionArg>>,
+ pub return_type: Option<DataType>,
+ // TODO: move this to new struct here
+ pub params: CreateFunctionBody,
+ //pub body: String,
+ /// Dummy schema
+ pub schema: DFSchemaRef,
+}
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct OperateFunctionArg {
+ // it is not really supported so no need to have it here
+ // currently
+ // pub mode: Option<ArgMode>,
+ pub name: Option<Ident>,
+ pub data_type: DataType,
+ pub default_expr: Option<Expr>,
+}
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunctionBody {
+ /// LANGUAGE lang_name
+ pub language: Option<Ident>,
+ /// IMMUTABLE | STABLE | VOLATILE
+ pub behavior: Option<Volatility>,
+ /// AS 'definition'
+ pub as_: Option<DefinitionStatement>,
+ /// RETURN expression
+ pub return_: Option<Expr>,
+}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub enum DefinitionStatement {
+ SingleQuotedDef(String),
+ DoubleDollarDef(String),
+}
+
+impl From<sqlparser::ast::FunctionDefinition> for DefinitionStatement {
+ fn from(value: sqlparser::ast::FunctionDefinition) -> Self {
+ match value {
+ sqlparser::ast::FunctionDefinition::SingleQuotedDef(s) => {
+ Self::SingleQuotedDef(s)
+ }
+ sqlparser::ast::FunctionDefinition::DoubleDollarDef(s) => {
+ Self::DoubleDollarDef(s)
+ }
+ }
+ }
+}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct DropFunction {
+ pub name: String,
+ pub if_exists: bool,
+ pub schema: DFSchemaRef,
Review Comment:
`DdlStatement::schema` needs it. I've tried to remove it and create it when needed but compiler complains it does not live long enough
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1975077302
> also, can you please let me know once new version of sqlparser lands, will remove dialect configuration from test.
I just released sqlparser https://crates.io/crates/sqlparser/0.44.0 (but we will need to update DataFusion to use it first)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1509957464
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -1261,7 +1313,38 @@ impl QueryPlanner for DefaultQueryPlanner {
.await
}
}
+/// 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 {
Review Comment:
I think the name is good and describes what this trait does -- namely it creates functions. I can't think of anything better or more appropriate
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -1306,6 +1381,12 @@ pub struct SessionState {
table_factories: HashMap<String, Arc<dyn TableProviderFactory>>,
/// Runtime environment
runtime_env: Arc<RuntimeEnv>,
+ /// [FunctionFactory] to support pluggable user defined function handler.
+ /// It is invoked on `CREATE FUNCTION` and `DROP FUNCTION` statements.
+ ///
+ /// Datafusion generic SQL dialect does not support `CRETE FUNCTION` statement
Review Comment:
BTW once I release the next version of sqlparser-rs (hopefully today) this will no longer be true
##########
datafusion/expr/src/logical_plan/ddl.rs:
##########
@@ -303,3 +321,66 @@ pub struct DropCatalogSchema {
/// Dummy schema
pub schema: DFSchemaRef,
}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunction {
+ // TODO: There is open question should we expose sqlparser types or redefine them here?
+ // At the moment it make more sense to expose sqlparser types and leave
+ // user to convert them as needed
+ pub or_replace: bool,
+ pub temporary: bool,
+ pub name: String,
+ pub args: Option<Vec<OperateFunctionArg>>,
+ pub return_type: Option<DataType>,
+ // TODO: move this to new struct here
+ pub params: CreateFunctionBody,
+ //pub body: String,
+ /// Dummy schema
+ pub schema: DFSchemaRef,
Review Comment:
yes, I meant I think we can remove the schema here
##########
datafusion/expr/src/logical_plan/ddl.rs:
##########
@@ -303,3 +321,66 @@ pub struct DropCatalogSchema {
/// Dummy schema
pub schema: DFSchemaRef,
}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunction {
+ // TODO: There is open question should we expose sqlparser types or redefine them here?
+ // At the moment it make more sense to expose sqlparser types and leave
+ // user to convert them as needed
+ pub or_replace: bool,
+ pub temporary: bool,
+ pub name: String,
+ pub args: Option<Vec<OperateFunctionArg>>,
+ pub return_type: Option<DataType>,
+ // TODO: move this to new struct here
+ pub params: CreateFunctionBody,
+ //pub body: String,
+ /// Dummy schema
+ pub schema: DFSchemaRef,
+}
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct OperateFunctionArg {
+ // it is not really supported so no need to have it here
+ // currently
+ // pub mode: Option<ArgMode>,
+ pub name: Option<Ident>,
+ pub data_type: DataType,
+ pub default_expr: Option<Expr>,
+}
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunctionBody {
+ /// LANGUAGE lang_name
+ pub language: Option<Ident>,
+ /// IMMUTABLE | STABLE | VOLATILE
+ pub behavior: Option<Volatility>,
+ /// AS 'definition'
+ pub as_: Option<DefinitionStatement>,
+ /// RETURN expression
+ pub return_: Option<Expr>,
+}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub enum DefinitionStatement {
+ SingleQuotedDef(String),
+ DoubleDollarDef(String),
+}
+
+impl From<sqlparser::ast::FunctionDefinition> for DefinitionStatement {
+ fn from(value: sqlparser::ast::FunctionDefinition) -> Self {
+ match value {
+ sqlparser::ast::FunctionDefinition::SingleQuotedDef(s) => {
+ Self::SingleQuotedDef(s)
+ }
+ sqlparser::ast::FunctionDefinition::DoubleDollarDef(s) => {
+ Self::DoubleDollarDef(s)
+ }
+ }
+ }
+}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct DropFunction {
+ pub name: String,
+ pub if_exists: bool,
+ pub schema: DFSchemaRef,
Review Comment:
Got it -- I tried locally and i see that `DdlStatement::schema()` returns a reference. Makes sense
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1510211503
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -794,6 +797,55 @@ impl SessionContext {
Ok(false)
}
+ async fn create_function(&self, stmt: CreateFunction) -> Result<DataFrame> {
+ let function = {
+ let state = self.state.read().clone();
+ let function_factory = &state.function_factory;
+
+ match function_factory {
+ Some(f) => f.create(state.config(), stmt).await?,
+ _ => Err(DataFusionError::Configuration(
+ "Function factory has not been configured".into(),
+ ))?,
+ }
+ };
+
+ match function {
+ RegisterFunction::Scalar(f) => {
+ self.state.write().register_udf(f)?;
+ }
+ RegisterFunction::Aggregate(f) => {
+ self.state.write().register_udaf(f)?;
+ }
+ RegisterFunction::Window(f) => {
+ self.state.write().register_udwf(f)?;
+ }
+ RegisterFunction::Table(name, f) => self.register_udtf(&name, f),
+ };
+
+ self.return_empty_dataframe()
+ }
+
+ async fn drop_function(&self, stmt: DropFunction) -> Result<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 {
Review Comment:
It is better to add the `if exist` test like `DROP FUNCTION if exists abs;` so the logic here is covered.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1509905879
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -619,6 +624,117 @@ async fn verify_udf_return_type() -> Result<()> {
Ok(())
}
+#[derive(Debug, Default)]
+struct MockFunctionFactory {
+ pub captured_expr: Mutex<Option<Expr>>,
+}
+
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ _config: &SessionConfig,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<RegisterFunction> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // capture expression so we can verify
+ // it has been parsed
+ *self.captured_expr.lock() = statement.params.return_;
+
+ Ok(RegisterFunction::Scalar(Arc::new(mock_udf)))
+ }
+
+ async fn remove(
Review Comment:
`remove` has been gone :), `DROP FUNCTION` will not discriminate against udf type and drop them all.
`remove` is dropped with condition that `simplify` from #9304 find its way to other function types, otherwise function factory would have to implement complicated registry with weak ref or something like that
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1509951359
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -635,6 +640,102 @@ async fn verify_udf_return_type() -> Result<()> {
Ok(())
}
+#[derive(Debug, Default)]
+struct MockFunctionFactory {
+ pub captured_expr: Mutex<Option<Expr>>,
+}
+
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ _config: &SessionConfig,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<RegisterFunction> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // capture expression so we can verify
+ // it has been parsed
+ *self.captured_expr.lock() = statement.params.return_;
+
+ Ok(RegisterFunction::Scalar(Arc::new(mock_udf)))
+ }
+}
+
+#[tokio::test]
+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();
+ let session_config =
+ SessionConfig::new().set_str("datafusion.sql_parser.dialect", "PostgreSQL");
+
+ let state =
+ SessionState::new_with_config_rt(session_config, Arc::new(runtime_environment))
+ .with_function_factory(function_factory.clone());
+
+ let ctx = SessionContext::new_with_state(state);
+
+ let sql = r#"
+ CREATE FUNCTION better_add(DOUBLE, DOUBLE)
+ RETURNS DOUBLE
+ RETURN $1 + $2
+ "#;
+ let _ = ctx.sql(sql).await?;
Review Comment:
Is this a comment to be added?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1979733384
Merged up from main again to resolve conflicts
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1975074784
> Its great to have a test with invalid function, the function which is not parseable or compilable.
>
> And the next question if the user created functions `f1` and `f2`, can `f1` be called inside `f2`?
I think it is important to note that this PR does NOT actually add any support for `CREATE FUNCTION` by default in DataFusion. Instead it sets up the plumbing for people to add their own `CREATE FUNCTION` support
Thus, the answer to this question is "it depends on whatever the user has coded the `FunctionFactory` to do
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1508163312
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -619,6 +624,117 @@ async fn verify_udf_return_type() -> Result<()> {
Ok(())
}
+#[derive(Debug, Default)]
+struct MockFunctionFactory {
+ pub captured_expr: Mutex<Option<Expr>>,
+}
+
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ _config: &SessionConfig,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<RegisterFunction> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // capture expression so we can verify
+ // it has been parsed
+ *self.captured_expr.lock() = statement.params.return_;
+
+ Ok(RegisterFunction::Scalar(Arc::new(mock_udf)))
+ }
+
+ async fn remove(
Review Comment:
Thinking about it more, maybe we don't even need a `remove` in the `FunctionFactory` at all -- dropping the function is mostly handled by `deregister_udf`?
Or maybe there is some sort of deiniitialization of the function that is needed that would be `FunctionFactory` specific๐ค. However, perhaps UDFs that needed that could implement some sort of deregistering / deallocation on `drop`.
Handling deregistering on `drop` might be better anyways as would prevents a function from being deregistered out from under a running execution...
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -794,6 +797,52 @@ impl SessionContext {
Ok(false)
}
+ async fn create_function(&self, stmt: CreateFunction) -> Result<DataFrame> {
+ let function = {
+ let state = self.state.read().clone();
+ let function_factory = &state.function_factory;
+
+ match function_factory {
+ Some(f) => f.create(state.config(), stmt).await?,
+ _ => Err(DataFusionError::Configuration(
+ "Function factory has not been configured".into(),
+ ))?,
+ }
+ };
+
+ match function {
+ RegisterFunction::Scalar(f) => {
+ self.state.write().register_udf(f)?;
+ }
+ RegisterFunction::Aggregate(f) => {
+ self.state.write().register_udaf(f)?;
+ }
+ RegisterFunction::Window(f) => {
+ self.state.write().register_udwf(f)?;
+ }
+ RegisterFunction::Table(name, f) => self.register_udtf(&name, f),
+ };
+
+ self.return_empty_dataframe()
+ }
+
+ async fn drop_function(&self, stmt: DropFunction) -> Result<DataFrame> {
+ 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
Review Comment:
I think we do have them (added relatively recently): https://github.com/apache/arrow-datafusion/blob/ea30b93194cfdfe4148a45fc0e33549884ba81b1/datafusion/core/src/execution/context/mod.rs#L853-L865
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1962442441
Drop function will be added as well
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1509360221
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -619,6 +624,117 @@ async fn verify_udf_return_type() -> Result<()> {
Ok(())
}
+#[derive(Debug, Default)]
+struct MockFunctionFactory {
+ pub captured_expr: Mutex<Option<Expr>>,
+}
+
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ _config: &SessionConfig,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<RegisterFunction> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // capture expression so we can verify
+ // it has been parsed
+ *self.captured_expr.lock() = statement.params.return_;
+
+ Ok(RegisterFunction::Scalar(Arc::new(mock_udf)))
+ }
+
+ async fn remove(
Review Comment:
I think that's good solution.
On that note there can be UDF and udaf with same name, if used in aggregation query, analyser will freak out, as it will find UDF function.
Will remove `remove` function from trait for now, we can always put it back
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1506768428
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -1246,7 +1273,42 @@ 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");
+/// ```
+#[async_trait]
+pub trait FunctionFactory: Sync + Send {
+ // TODO: I don't like having RwLock Leaking here, who ever implements it
+ // has to depend ot `parking_lot`. I'f we expose &mut SessionState it
+ // may keep lock of too long.
+ //
+ // Not sure if there is better approach.
+ //
+
+ /// Handles creation of user defined function specified in [CreateFunction] statement
+ async fn create(
+ &self,
+ state: Arc<RwLock<SessionState>>,
Review Comment:
Should we want to implement sql user defined functions, we would need extension to cache logical expressions and provide them to logical analyser so it can replace udf call with the expression.
Id argue we need state or configuration as a parameter as well. Let me do other changes, as they are straightforward and we can try to find most ergonomic function
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1508153276
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -794,6 +797,52 @@ impl SessionContext {
Ok(false)
}
+ async fn create_function(&self, stmt: CreateFunction) -> Result<DataFrame> {
Review Comment:
๐
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -1246,7 +1295,51 @@ 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`
Review Comment:
I agree supporting create function would be good in the generic dialect. I made this change in https://github.com/sqlparser-rs/sqlparser-rs/pull/1159 (and I plan a sqlparser release in the next few days)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1974876415
> > Thanks @milenkovicm I would say it is nice to get also tests what will happen if the function is invalid/not compilable or something
> > another question if function can call another function?
>
> Thanks for your comments @comphead , can you please elaborate your comment, thanks
Its great to have a test with invalid function, the function which is not parseable or compilable.
And the next question if the user created functions `f1` and `f2`, can `f1` be called inside `f2`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1975077103
I took the liberty of pushing a commit with a sqllogictest that shows the behavior of `CREATE FUNCTION` / `DROP FUNCTION` which should hopefully make this PR a little clearer
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1977568695
https://github.com/apache/arrow-datafusion/pull/9268 was merged, so the new sqlparser release should now be available.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1508200064
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -619,6 +624,117 @@ async fn verify_udf_return_type() -> Result<()> {
Ok(())
}
+#[derive(Debug, Default)]
+struct MockFunctionFactory {
+ pub captured_expr: Mutex<Option<Expr>>,
+}
+
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ _config: &SessionConfig,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<RegisterFunction> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // capture expression so we can verify
+ // it has been parsed
+ *self.captured_expr.lock() = statement.params.return_;
+
+ Ok(RegisterFunction::Scalar(Arc::new(mock_udf)))
+ }
+
+ async fn remove(
Review Comment:
> do you maybe have idea how could we decide which unregister_ud*f to call when we get DROP FUNCTION?
That is a (very) good question. I don't have a good answer at the moment. I'll try and think of one
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1974878093
> > > Thanks @milenkovicm I would say it is nice to get also tests what will happen if the function is invalid/not compilable or something
> > > another question if function can call another function?
> >
> >
> > Thanks for your comments @comphead , can you please elaborate your comment, thanks
>
> Its great to have a test with invalid function, the function which is not parseable or compilable.
could be added, no problem
> And the next question if the user created functions `f1` and `f2`, can `f1` be called inside `f2`?
I see two cases here:
1. if we have functions defined in rust, there is no change to current behaviour I believe
2. if we have `SQL UDF` then analyser should replace f1 first, then f2, so it should run as well
do i miss something here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1510017098
##########
datafusion/core/tests/sql/sql_api.rs:
##########
@@ -39,6 +52,99 @@ async fn unsupported_ddl_returns_error() {
ctx.sql_with_options(sql, options).await.unwrap();
}
+struct MockFunctionFactory;
Review Comment:
Right, this functionality looks important and definitely needs to be covered with docs, examples in following PR, looking forward for it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1979704066
I merged up from main to pick up a fix for the CI
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1976342371
> Please let me know when new sqlparser lands so I can update dialect used in test
Thank you -- I saw https://github.com/apache/arrow-datafusion/pull/9447 from dependabot and I think @JasonLi-cn has implemented the upgrade as part of https://github.com/apache/arrow-datafusion/pull/9268
I will attempt to review https://github.com/apache/arrow-datafusion/pull/9268 later today
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1509951559
##########
datafusion/expr/src/logical_plan/ddl.rs:
##########
@@ -303,3 +321,66 @@ pub struct DropCatalogSchema {
/// Dummy schema
pub schema: DFSchemaRef,
}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunction {
+ // TODO: There is open question should we expose sqlparser types or redefine them here?
+ // At the moment it make more sense to expose sqlparser types and leave
+ // user to convert them as needed
+ pub or_replace: bool,
+ pub temporary: bool,
+ pub name: String,
+ pub args: Option<Vec<OperateFunctionArg>>,
+ pub return_type: Option<DataType>,
+ // TODO: move this to new struct here
+ pub params: CreateFunctionBody,
+ //pub body: String,
+ /// Dummy schema
+ pub schema: DFSchemaRef,
Review Comment:
Its not needed
##########
datafusion/expr/src/logical_plan/ddl.rs:
##########
@@ -303,3 +321,66 @@ pub struct DropCatalogSchema {
/// Dummy schema
pub schema: DFSchemaRef,
}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunction {
+ // TODO: There is open question should we expose sqlparser types or redefine them here?
+ // At the moment it make more sense to expose sqlparser types and leave
+ // user to convert them as needed
+ pub or_replace: bool,
+ pub temporary: bool,
+ pub name: String,
+ pub args: Option<Vec<OperateFunctionArg>>,
+ pub return_type: Option<DataType>,
+ // TODO: move this to new struct here
+ pub params: CreateFunctionBody,
+ //pub body: String,
+ /// Dummy schema
+ pub schema: DFSchemaRef,
Review Comment:
Its not needed, my fault
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1975086619
Thanks a lot @alamb
As follow up, maybe we should consider modifying create function syntax, to add support for named parameters and global variables ( if that's correct name `@var` and `@@var` )
Then next step would be adding `sql function contrib` ๐ similar to databrick's.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1507794956
##########
datafusion/core/tests/sql/sql_api.rs:
##########
@@ -39,6 +52,99 @@ async fn unsupported_ddl_returns_error() {
ctx.sql_with_options(sql, options).await.unwrap();
}
+struct MockFunctionFactory;
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ state: Arc<RwLock<SessionState>>,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<()> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // we may need other infrastructure provided by state, for example:
+ // state.config().get_extension()
Review Comment:
This comment is connected to previous, as getting `config` involves holding read lock. Read may not be as bad as write lock but still it will block writers.
As `create` may take some time (like getting a ML model from repo) read lock may be problematic as well.
Function signature definitely looks better
```rust
async fn create(
&self,
state: &SessionConfig,
statement: CreateFunction,
) -> Result<RegisterFunction>;
```
but lock is still an issue.
looks like `TableProviderFactory` has same signature:
```rust
#[async_trait]
pub trait TableProviderFactory: Sync + Send {
/// Create a TableProvider with the given url
async fn create(
&self,
state: &SessionState,
cmd: &CreateExternalTable,
) -> Result<Arc<dyn TableProvider>>;
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1507799073
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -1246,7 +1273,42 @@ 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");
+/// ```
+#[async_trait]
+pub trait FunctionFactory: Sync + Send {
+ // TODO: I don't like having RwLock Leaking here, who ever implements it
+ // has to depend ot `parking_lot`. I'f we expose &mut SessionState it
+ // may keep lock of too long.
+ //
+ // Not sure if there is better approach.
+ //
+
+ /// Handles creation of user defined function specified in [CreateFunction] statement
+ async fn create(
+ &self,
+ state: Arc<RwLock<SessionState>>,
Review Comment:
I definitely like this approach, but it has lock problem if we want to provide a `config` parameter.
Why would we need `config`? `FunctionFactory` may store configuration for model repository url, or similar
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1508171908
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -619,6 +624,117 @@ async fn verify_udf_return_type() -> Result<()> {
Ok(())
}
+#[derive(Debug, Default)]
+struct MockFunctionFactory {
+ pub captured_expr: Mutex<Option<Expr>>,
+}
+
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ _config: &SessionConfig,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<RegisterFunction> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // capture expression so we can verify
+ // it has been parsed
+ *self.captured_expr.lock() = statement.params.return_;
+
+ Ok(RegisterFunction::Scalar(Arc::new(mock_udf)))
+ }
+
+ async fn remove(
Review Comment:
lets say we have a SQL UDFs which are expressed by expr, we would need to have some kind of registry (which is not udf_registry) which can be consulted by plan analyser.
so in this case we would have some kind of mock function stored in function registry (which would do nothing, just for planning) and have an expression registry, which would be used to "replace" udf's in logical plan ...
btw, i believe adding sql udf would be trivial once we get this change and fix few issues with expression namings
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1974785906
> Thanks again @milenkovicm -- the only thing left from my perspective is to "error on multi-part identifiers" which I think we could do as a follow on PR if needed.
>
> However, I do think we should leave this open for a few more days to let other people have a chance to comment on it before we merge.
>
> Thanks again. This is a _really_ neat feature that I think will let people build all sorts of cool things using DataFusion ๐
@alamb multi-part identifiers fixed as well, please have a look.
also, can you please let me know once new version of `sqlparser` lands, will remove dialect configuration from test.
thanks a lot for all your help
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1509959181
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -1306,6 +1381,12 @@ pub struct SessionState {
table_factories: HashMap<String, Arc<dyn TableProviderFactory>>,
/// Runtime environment
runtime_env: Arc<RuntimeEnv>,
+ /// [FunctionFactory] to support pluggable user defined function handler.
+ /// It is invoked on `CREATE FUNCTION` and `DROP FUNCTION` statements.
+ ///
+ /// Datafusion generic SQL dialect does not support `CRETE FUNCTION` statement
Review Comment:
updated
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1974832950
> and please squash commit on merge, thanks
Yes, this will be done (as on all commits)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1508167240
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -1246,7 +1295,51 @@ 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`
Review Comment:
I think we added the unregister_* functions recently (sorry our reviews crossed)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1508164210
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -1246,7 +1295,51 @@ 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`
Review Comment:
looking at `FunctionRegistry.remove` method, in order to implement it `FunctionRegistry` should be changed. Maybe that change is a bit too much for this PR. So i propose to remove `remove` from `FunctionFactory` and add it back when we implement `unregister_ud*` methods in `FunctionRegistry`.
wdyt?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1976335180
I plan to leave this PR open for another day or two to give people a chance to comment if they would like to and then I'll merge it in.
Thanks again @milenkovicm
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1507794956
##########
datafusion/core/tests/sql/sql_api.rs:
##########
@@ -39,6 +52,99 @@ async fn unsupported_ddl_returns_error() {
ctx.sql_with_options(sql, options).await.unwrap();
}
+struct MockFunctionFactory;
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ state: Arc<RwLock<SessionState>>,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<()> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // we may need other infrastructure provided by state, for example:
+ // state.config().get_extension()
Review Comment:
This comment is connected to previous, as getting `config` involves holding read lock. Read may not be as bad as write lock but still it will block writers.
As `create` may take some time (like getting a ML model from repo) read lock may be problematic as well.
Function signature definitely looks better
```rust
async fn create(
&self,
state: &SessionConfig,
statement: CreateFunction,
) -> Result<RegisterFunction>;
```
but lock is still an issue.
looks like `TableProviderFactory` has same signature:
```rust
#[async_trait]
pub trait TableProviderFactory: Sync + Send {
/// Create a TableProvider with the given url
async fn create(
&self,
state: &SessionState,
cmd: &CreateExternalTable,
) -> Result<Arc<dyn TableProvider>>;
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1508156947
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -1246,7 +1273,42 @@ 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");
+/// ```
+#[async_trait]
+pub trait FunctionFactory: Sync + Send {
+ // TODO: I don't like having RwLock Leaking here, who ever implements it
+ // has to depend ot `parking_lot`. I'f we expose &mut SessionState it
+ // may keep lock of too long.
+ //
+ // Not sure if there is better approach.
+ //
+
+ /// Handles creation of user defined function specified in [CreateFunction] statement
+ async fn create(
+ &self,
+ state: Arc<RwLock<SessionState>>,
Review Comment:
Sorry for confusion @alamb, I believe i fix it with right amount of .clone() on right places :).
Please have a look now, I believe it looks all right. I still do not like `remove(...)` could be implemented with another enum ... have o think about it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1510016448
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -794,6 +797,55 @@ impl SessionContext {
Ok(false)
}
+ async fn create_function(&self, stmt: CreateFunction) -> Result<DataFrame> {
+ let function = {
+ let state = self.state.read().clone();
+ let function_factory = &state.function_factory;
+
+ match function_factory {
+ Some(f) => f.create(state.config(), stmt).await?,
+ _ => Err(DataFusionError::Configuration(
+ "Function factory has not been configured".into(),
+ ))?,
+ }
+ };
+
+ match function {
+ RegisterFunction::Scalar(f) => {
+ self.state.write().register_udf(f)?;
+ }
+ RegisterFunction::Aggregate(f) => {
+ self.state.write().register_udaf(f)?;
+ }
+ RegisterFunction::Window(f) => {
+ self.state.write().register_udwf(f)?;
+ }
+ RegisterFunction::Table(name, f) => self.register_udtf(&name, f),
+ };
+
+ self.return_empty_dataframe()
+ }
+
+ async fn drop_function(&self, stmt: DropFunction) -> Result<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()))
Review Comment:
please use `exec_err!` macros
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1506694977
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -495,6 +496,32 @@ impl SessionContext {
self.set_variable(stmt).await
}
+ LogicalPlan::Statement(Statement::CreateFunction(stmt)) => {
Review Comment:
Since CreateFunction / DropFunction modify the catalog, what do you think about making them part of `DdlStatement` instead of a `Statement`?
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -495,6 +496,32 @@ impl SessionContext {
self.set_variable(stmt).await
}
+ LogicalPlan::Statement(Statement::CreateFunction(stmt)) => {
+ let function_factory = self.state.read().function_factory.clone();
+
+ match function_factory {
+ Some(f) => f.create(self.state.clone(), stmt).await?,
+ None => Err(DataFusionError::Configuration(
+ "Function factory has not been configured".into(),
+ ))?,
+ };
Review Comment:
Minor: I think it would be better to follow the model of `create_view` etc and make this code a function on `SessionContext` (so that this code just dispatches to that method)
The upside of doing so is that it follows the existing code pattern and it offers a place to document `create_function` and `drop_function` publically in in more detail
##########
datafusion/core/tests/sql/sql_api.rs:
##########
@@ -39,6 +52,99 @@ async fn unsupported_ddl_returns_error() {
ctx.sql_with_options(sql, options).await.unwrap();
}
+struct MockFunctionFactory;
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ state: Arc<RwLock<SessionState>>,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<()> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // we may need other infrastructure provided by state, for example:
+ // state.config().get_extension()
Review Comment:
๐ค maybe we should pass the config into `create` rather than `SessionContext`
##########
datafusion/expr/src/logical_plan/statement.rs:
##########
@@ -148,3 +164,32 @@ pub struct SetVariable {
/// Dummy schema
pub schema: DFSchemaRef,
}
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunction {
+ // TODO: There is open question should we expose sqlparser types or redefine them here?
Review Comment:
I think that this approach (redefine the types) is more consistent with the rest of the DataFusion codebase, though it is an unfortunate amount of replication
##########
datafusion/core/tests/sql/sql_api.rs:
##########
@@ -39,6 +52,99 @@ async fn unsupported_ddl_returns_error() {
ctx.sql_with_options(sql, options).await.unwrap();
}
+struct MockFunctionFactory;
Review Comment:
This is very cool -- I think it would be great (as a follow on PR) if we turned this into an example (with docs, etc like the others in https://github.com/apache/arrow-datafusion/tree/main/datafusion-examples/examples)
##########
datafusion/core/tests/sql/sql_api.rs:
##########
@@ -39,6 +52,99 @@ async fn unsupported_ddl_returns_error() {
ctx.sql_with_options(sql, options).await.unwrap();
}
+struct MockFunctionFactory;
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ state: Arc<RwLock<SessionState>>,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<()> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // we may need other infrastructure provided by state, for example:
+ // state.config().get_extension()
+
+ // register mock udf for testing
+ state.write().register_udf(mock_udf.into())?;
+ Ok(())
+ }
+
+ async fn remove(
+ &self,
+ _state: Arc<RwLock<SessionState>>,
+ _statement: DropFunction,
+ ) -> datafusion::error::Result<()> {
+ // at the moment state does not support unregister
+ // ignoring for now
+ Ok(())
+ }
+}
+
+#[tokio::test]
+async fn create_user_defined_function_statement() {
Review Comment:
Perhaps we could put this test along side the other user defined tests in https://github.com/apache/arrow-datafusion/tree/main/datafusion/core/tests/user_defined so it was easier to find
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -1246,7 +1273,42 @@ 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");
+/// ```
+#[async_trait]
+pub trait FunctionFactory: Sync + Send {
+ // TODO: I don't like having RwLock Leaking here, who ever implements it
+ // has to depend ot `parking_lot`. I'f we expose &mut SessionState it
+ // may keep lock of too long.
+ //
+ // Not sure if there is better approach.
+ //
+
+ /// Handles creation of user defined function specified in [CreateFunction] statement
+ async fn create(
+ &self,
+ state: Arc<RwLock<SessionState>>,
Review Comment:
It seems like `create` here handles two things:
1. Creation of the function
2. Registering the function with the SessionState
What if we made `create` only responsible for creating a function and made the caller responsible for registering it? I think that would make the API simpler / not involve RW lock
For example, something like
```rust
/// A function defined by a FunctionFactory
enum NewFunction {
Scalar(ScalarUDF),
Aggregate(AggregateUDF),
Window(WindowUDF)
}
pub trait FunctionFactory: Sync + Send {
...
async fn create(
&self,
statement: CreateFunction,
) -> Result<NewFunction>;
...
}
```
Another alternative might be rather than passing the `SessionState`, maybe we could pass a reference to `&mut dyn FunctionRegistry` which would also avoid the RwLock ๐ค
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1509953513
##########
datafusion/expr/src/logical_plan/ddl.rs:
##########
@@ -303,3 +321,66 @@ pub struct DropCatalogSchema {
/// Dummy schema
pub schema: DFSchemaRef,
}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunction {
+ // TODO: There is open question should we expose sqlparser types or redefine them here?
+ // At the moment it make more sense to expose sqlparser types and leave
+ // user to convert them as needed
+ pub or_replace: bool,
+ pub temporary: bool,
+ pub name: String,
+ pub args: Option<Vec<OperateFunctionArg>>,
+ pub return_type: Option<DataType>,
+ // TODO: move this to new struct here
+ pub params: CreateFunctionBody,
+ //pub body: String,
+ /// Dummy schema
+ pub schema: DFSchemaRef,
+}
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct OperateFunctionArg {
+ // it is not really supported so no need to have it here
+ // currently
+ // pub mode: Option<ArgMode>,
+ pub name: Option<Ident>,
+ pub data_type: DataType,
+ pub default_expr: Option<Expr>,
+}
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunctionBody {
+ /// LANGUAGE lang_name
+ pub language: Option<Ident>,
+ /// IMMUTABLE | STABLE | VOLATILE
+ pub behavior: Option<Volatility>,
+ /// AS 'definition'
+ pub as_: Option<DefinitionStatement>,
+ /// RETURN expression
+ pub return_: Option<Expr>,
+}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub enum DefinitionStatement {
+ SingleQuotedDef(String),
+ DoubleDollarDef(String),
+}
+
+impl From<sqlparser::ast::FunctionDefinition> for DefinitionStatement {
+ fn from(value: sqlparser::ast::FunctionDefinition) -> Self {
+ match value {
+ sqlparser::ast::FunctionDefinition::SingleQuotedDef(s) => {
+ Self::SingleQuotedDef(s)
+ }
+ sqlparser::ast::FunctionDefinition::DoubleDollarDef(s) => {
+ Self::DoubleDollarDef(s)
+ }
+ }
+ }
+}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct DropFunction {
+ pub name: String,
+ pub if_exists: bool,
+ pub schema: DFSchemaRef,
Review Comment:
Has to be `pub`. Other `structs` from that package have same property
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1508152349
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -1246,7 +1273,42 @@ 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");
+/// ```
+#[async_trait]
+pub trait FunctionFactory: Sync + Send {
+ // TODO: I don't like having RwLock Leaking here, who ever implements it
+ // has to depend ot `parking_lot`. I'f we expose &mut SessionState it
+ // may keep lock of too long.
+ //
+ // Not sure if there is better approach.
+ //
+
+ /// Handles creation of user defined function specified in [CreateFunction] statement
+ async fn create(
+ &self,
+ state: Arc<RwLock<SessionState>>,
Review Comment:
I think you are right. I mentioned `config` because I think there was a comment about it before.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1509329295
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -619,6 +624,117 @@ async fn verify_udf_return_type() -> Result<()> {
Ok(())
}
+#[derive(Debug, Default)]
+struct MockFunctionFactory {
+ pub captured_expr: Mutex<Option<Expr>>,
+}
+
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ _config: &SessionConfig,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<RegisterFunction> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // capture expression so we can verify
+ // it has been parsed
+ *self.captured_expr.lock() = statement.params.return_;
+
+ Ok(RegisterFunction::Scalar(Arc::new(mock_udf)))
+ }
+
+ async fn remove(
Review Comment:
I thought more about this last night and one potential solution is to simply try and deregister all functions with the specified name (udf, udaf, udwf, etc)
It seems non ideal however as it would allow dropping any udf function, not just the ones created by `CREATE FUNCTION` ๐ค
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1509945936
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -635,6 +640,102 @@ async fn verify_udf_return_type() -> Result<()> {
Ok(())
}
+#[derive(Debug, Default)]
+struct MockFunctionFactory {
+ pub captured_expr: Mutex<Option<Expr>>,
+}
+
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ _config: &SessionConfig,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<RegisterFunction> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // capture expression so we can verify
+ // it has been parsed
+ *self.captured_expr.lock() = statement.params.return_;
+
+ Ok(RegisterFunction::Scalar(Arc::new(mock_udf)))
+ }
+}
+
+#[tokio::test]
+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();
+ let session_config =
+ SessionConfig::new().set_str("datafusion.sql_parser.dialect", "PostgreSQL");
+
+ let state =
+ SessionState::new_with_config_rt(session_config, Arc::new(runtime_environment))
+ .with_function_factory(function_factory.clone());
+
+ let ctx = SessionContext::new_with_state(state);
+
+ let sql = r#"
+ CREATE FUNCTION better_add(DOUBLE, DOUBLE)
+ RETURNS DOUBLE
+ RETURN $1 + $2
+ "#;
+ let _ = ctx.sql(sql).await?;
+
+ 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();
+
+ // is there some better way to test this
Review Comment:
I think this is fine
##########
datafusion/expr/src/logical_plan/ddl.rs:
##########
@@ -303,3 +321,66 @@ pub struct DropCatalogSchema {
/// Dummy schema
pub schema: DFSchemaRef,
}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunction {
+ // TODO: There is open question should we expose sqlparser types or redefine them here?
+ // At the moment it make more sense to expose sqlparser types and leave
+ // user to convert them as needed
+ pub or_replace: bool,
+ pub temporary: bool,
+ pub name: String,
+ pub args: Option<Vec<OperateFunctionArg>>,
+ pub return_type: Option<DataType>,
+ // TODO: move this to new struct here
+ pub params: CreateFunctionBody,
+ //pub body: String,
+ /// Dummy schema
+ pub schema: DFSchemaRef,
Review Comment:
Why is this needed?
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -635,6 +640,102 @@ async fn verify_udf_return_type() -> Result<()> {
Ok(())
}
+#[derive(Debug, Default)]
+struct MockFunctionFactory {
+ pub captured_expr: Mutex<Option<Expr>>,
+}
+
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ _config: &SessionConfig,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<RegisterFunction> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
Review Comment:
```suggestion
// In this example, we always create a function that adds its arguments
// with the name specified in `CREATE FUNCTION`. In a real implementation
// the body of the created UDF would also likely be a function of the contents
// of the `CreateFunction`
```
##########
datafusion/expr/src/logical_plan/ddl.rs:
##########
@@ -303,3 +321,66 @@ pub struct DropCatalogSchema {
/// Dummy schema
pub schema: DFSchemaRef,
}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunction {
+ // TODO: There is open question should we expose sqlparser types or redefine them here?
+ // At the moment it make more sense to expose sqlparser types and leave
+ // user to convert them as needed
+ pub or_replace: bool,
+ pub temporary: bool,
+ pub name: String,
+ pub args: Option<Vec<OperateFunctionArg>>,
+ pub return_type: Option<DataType>,
+ // TODO: move this to new struct here
+ pub params: CreateFunctionBody,
+ //pub body: String,
+ /// Dummy schema
+ pub schema: DFSchemaRef,
+}
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct OperateFunctionArg {
+ // it is not really supported so no need to have it here
+ // currently
+ // pub mode: Option<ArgMode>,
+ pub name: Option<Ident>,
+ pub data_type: DataType,
+ pub default_expr: Option<Expr>,
+}
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunctionBody {
+ /// LANGUAGE lang_name
+ pub language: Option<Ident>,
+ /// IMMUTABLE | STABLE | VOLATILE
+ pub behavior: Option<Volatility>,
+ /// AS 'definition'
+ pub as_: Option<DefinitionStatement>,
+ /// RETURN expression
+ pub return_: Option<Expr>,
+}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub enum DefinitionStatement {
+ SingleQuotedDef(String),
+ DoubleDollarDef(String),
+}
+
+impl From<sqlparser::ast::FunctionDefinition> for DefinitionStatement {
+ fn from(value: sqlparser::ast::FunctionDefinition) -> Self {
+ match value {
+ sqlparser::ast::FunctionDefinition::SingleQuotedDef(s) => {
+ Self::SingleQuotedDef(s)
+ }
+ sqlparser::ast::FunctionDefinition::DoubleDollarDef(s) => {
+ Self::DoubleDollarDef(s)
+ }
+ }
+ }
+}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct DropFunction {
+ pub name: String,
+ pub if_exists: bool,
+ pub schema: DFSchemaRef,
Review Comment:
Does this really need a schema?
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -1568,6 +1650,16 @@ impl SessionState {
self
}
+ /// Registers `CREATE FUNCTION` statement handler implementing
+ /// [`FunctionFactory`] trait.
Review Comment:
```suggestion
/// Registers a [`FunctionFactory`] to handle `CREATE FUNCTION` statements
```
##########
datafusion/sql/src/statement.rs:
##########
@@ -625,8 +626,136 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
});
Ok(LogicalPlan::Statement(statement))
}
+ Statement::CreateFunction {
+ or_replace,
+ temporary,
+ name,
+ args,
+ return_type,
+ params,
+ } => {
+ let return_type = match return_type {
+ Some(t) => Some(self.convert_data_type(&t)?),
+ None => None,
+ };
+ let mut planner_context = PlannerContext::new();
+ let empty_schema = &DFSchema::empty();
+
+ let args = match args {
+ Some(function_args) => {
+ let function_args = function_args
+ .into_iter()
+ .map(|arg| {
+ let data_type = self.convert_data_type(&arg.data_type)?;
+
+ let default_expr = match arg.default_expr {
+ Some(expr) => Some(self.sql_to_expr(
+ expr,
+ empty_schema,
+ &mut planner_context,
+ )?),
+ None => None,
+ };
+ Ok(OperateFunctionArg {
+ name: arg.name,
+ default_expr,
+ data_type,
+ })
+ })
+ .collect::<Result<Vec<OperateFunctionArg>>>();
+ Some(function_args?)
+ }
+ None => None,
+ };
+
+ // Not sure if this is correct way to generate name
+ // postgresql function definition may have schema part as well
+ // datafusion at the moment does lookup based on given string
+ // `schema_name.function_name` will work even if there is no `schema_name`
+ let name: String = name
+ .0
+ .into_iter()
+ .map(|i| i.value)
+ .collect::<Vec<String>>()
+ .join(".");
Review Comment:
I recommend in this case that we throw an error if someone tries to create a qualified function, and if it is important to support qualified function names in the future we can revisit the issue then
So
```
CREATE FUNCTION foo
```
would work
but
```
CREATE FUNCTION my.foo ...
```
would error (not yet implemented error)
This is the corresponding code for table references https://github.com/apache/arrow-datafusion/blob/dd4263f843e093c807d63edf73a571b1ba2669b5/datafusion/common/src/column.rs#L71-L99
is https://docs.rs/datafusion/latest/datafusion/sql/enum.TableReference.html
But at the moment DataFusion doesn't put functions in a schema or catalog, it simply puts them some global namespace.
##########
datafusion/expr/src/logical_plan/ddl.rs:
##########
@@ -303,3 +321,66 @@ pub struct DropCatalogSchema {
/// Dummy schema
pub schema: DFSchemaRef,
}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunction {
Review Comment:
```suggestion
/// Arguments passed to `CREATE FUNCTION`
///
/// Note this meant to be the same as from sqlparser's [`sqlparser::ast::Statement::CreateFunction`]
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct CreateFunction {
```
##########
datafusion/expr/src/logical_plan/ddl.rs:
##########
@@ -303,3 +321,66 @@ pub struct DropCatalogSchema {
/// Dummy schema
pub schema: DFSchemaRef,
}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunction {
+ // TODO: There is open question should we expose sqlparser types or redefine them here?
Review Comment:
Is this still TODO (or are the rest of the comments below todo)? I think this structure makes sense and with a little more explanation on rationale (aka the Statement::CreateFunction is a variant and not its own struct) this make sense
The rationale I think would be important so that when we upgrade sqlparser's version and it adds more CREATE FUNCTION variants, the updater knows the desired results in DataFusion
##########
datafusion/sql/src/statement.rs:
##########
@@ -625,8 +626,136 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
});
Ok(LogicalPlan::Statement(statement))
}
+ Statement::CreateFunction {
+ or_replace,
+ temporary,
+ name,
+ args,
+ return_type,
+ params,
+ } => {
+ let return_type = match return_type {
+ Some(t) => Some(self.convert_data_type(&t)?),
+ None => None,
+ };
+ let mut planner_context = PlannerContext::new();
+ let empty_schema = &DFSchema::empty();
+
+ let args = match args {
+ Some(function_args) => {
+ let function_args = function_args
+ .into_iter()
+ .map(|arg| {
+ let data_type = self.convert_data_type(&arg.data_type)?;
+
+ let default_expr = match arg.default_expr {
+ Some(expr) => Some(self.sql_to_expr(
+ expr,
+ empty_schema,
+ &mut planner_context,
+ )?),
+ None => None,
+ };
+ Ok(OperateFunctionArg {
+ name: arg.name,
+ default_expr,
+ data_type,
+ })
+ })
+ .collect::<Result<Vec<OperateFunctionArg>>>();
+ Some(function_args?)
+ }
+ None => None,
+ };
+
+ // Not sure if this is correct way to generate name
+ // postgresql function definition may have schema part as well
+ // datafusion at the moment does lookup based on given string
+ // `schema_name.function_name` will work even if there is no `schema_name`
+ let name: String = name
+ .0
+ .into_iter()
+ .map(|i| i.value)
+ .collect::<Vec<String>>()
+ .join(".");
+
+ //
+ // convert resulting expression to data fusion expression
+ //
+ let arg_types = args.as_ref().map(|arg| {
+ arg.iter().map(|t| t.data_type.clone()).collect::<Vec<_>>()
+ });
+ let mut planner_context = PlannerContext::new()
+ .with_prepare_param_data_types(arg_types.unwrap_or_default());
+
+ let result_expression = match params.return_ {
+ Some(r) => Some(self.sql_to_expr(
+ r,
+ &DFSchema::empty(),
+ &mut planner_context,
+ )?),
+ None => None,
+ };
+
+ let params = CreateFunctionBody {
+ language: params.language,
+ behavior: params.behavior.map(|b| match b {
+ ast::FunctionBehavior::Immutable => Volatility::Immutable,
+ ast::FunctionBehavior::Stable => Volatility::Stable,
+ ast::FunctionBehavior::Volatile => Volatility::Volatile,
+ }),
+ as_: params.as_.map(|m| m.into()),
+ return_: result_expression,
+ };
- _ => not_impl_err!("Unsupported SQL statement: {sql:?}"),
+ let statement = DdlStatement::CreateFunction(CreateFunction {
+ or_replace,
+ temporary,
+ name,
+ return_type,
+ args,
+ params,
+ schema: DFSchemaRef::new(DFSchema::empty()),
+ });
+
+ Ok(LogicalPlan::Ddl(statement))
+ }
+ Statement::DropFunction {
+ if_exists,
+ func_desc,
+ ..
+ } => {
+ // Not sure if this is correct way to generate name
Review Comment:
See above -- I also recommend erroring in the case of multi part identifiers.
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -635,6 +640,102 @@ async fn verify_udf_return_type() -> Result<()> {
Ok(())
}
+#[derive(Debug, Default)]
+struct MockFunctionFactory {
+ pub captured_expr: Mutex<Option<Expr>>,
+}
+
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ _config: &SessionConfig,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<RegisterFunction> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // capture expression so we can verify
+ // it has been parsed
+ *self.captured_expr.lock() = statement.params.return_;
+
+ Ok(RegisterFunction::Scalar(Arc::new(mock_udf)))
+ }
+}
+
+#[tokio::test]
+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();
+ let session_config =
+ SessionConfig::new().set_str("datafusion.sql_parser.dialect", "PostgreSQL");
+
+ let state =
+ SessionState::new_with_config_rt(session_config, Arc::new(runtime_environment))
+ .with_function_factory(function_factory.clone());
+
+ let ctx = SessionContext::new_with_state(state);
+
+ let sql = r#"
+ CREATE FUNCTION better_add(DOUBLE, DOUBLE)
+ RETURNS DOUBLE
+ RETURN $1 + $2
+ "#;
+ let _ = ctx.sql(sql).await?;
Review Comment:
Create the `better_add` function dynamically via `CREATE FUNCTION` statement
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -635,6 +640,102 @@ async fn verify_udf_return_type() -> Result<()> {
Ok(())
}
+#[derive(Debug, Default)]
+struct MockFunctionFactory {
+ pub captured_expr: Mutex<Option<Expr>>,
+}
+
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ _config: &SessionConfig,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<RegisterFunction> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // capture expression so we can verify
+ // it has been parsed
+ *self.captured_expr.lock() = statement.params.return_;
+
+ Ok(RegisterFunction::Scalar(Arc::new(mock_udf)))
+ }
+}
+
+#[tokio::test]
+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();
+ let session_config =
+ SessionConfig::new().set_str("datafusion.sql_parser.dialect", "PostgreSQL");
+
+ let state =
+ SessionState::new_with_config_rt(session_config, Arc::new(runtime_environment))
+ .with_function_factory(function_factory.clone());
+
+ let ctx = SessionContext::new_with_state(state);
+
+ let sql = r#"
+ CREATE FUNCTION better_add(DOUBLE, DOUBLE)
+ RETURNS DOUBLE
+ RETURN $1 + $2
+ "#;
+ let _ = ctx.sql(sql).await?;
+
+ 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();
+
+ // is there some better way to test this
+ assert_eq!("$1 + $2", captured_expression.to_string());
+
+ // 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());
+
Review Comment:
Would it also be possible to add a test using `SessionContext::sql_with_options` showing that if DDL is disabled that CREATE / DROP function is also disabled?
https://docs.rs/datafusion/latest/datafusion/execution/context/struct.SessionContext.html#method.sql_with_options
##########
datafusion/expr/src/logical_plan/ddl.rs:
##########
@@ -303,3 +321,66 @@ pub struct DropCatalogSchema {
/// Dummy schema
pub schema: DFSchemaRef,
}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunction {
+ // TODO: There is open question should we expose sqlparser types or redefine them here?
+ // At the moment it make more sense to expose sqlparser types and leave
+ // user to convert them as needed
+ pub or_replace: bool,
+ pub temporary: bool,
+ pub name: String,
+ pub args: Option<Vec<OperateFunctionArg>>,
+ pub return_type: Option<DataType>,
+ // TODO: move this to new struct here
+ pub params: CreateFunctionBody,
+ //pub body: String,
+ /// Dummy schema
+ pub schema: DFSchemaRef,
+}
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct OperateFunctionArg {
+ // it is not really supported so no need to have it here
+ // currently
Review Comment:
I think it is fine to leave some of these fields as "TODO" items and if someone needs them they can add them as a follow on PR
```suggestion
// TODO: figure out how to support mode
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1968868488
Hi @milenkovicm -- I plan to review this PR shortly
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1507794956
##########
datafusion/core/tests/sql/sql_api.rs:
##########
@@ -39,6 +52,99 @@ async fn unsupported_ddl_returns_error() {
ctx.sql_with_options(sql, options).await.unwrap();
}
+struct MockFunctionFactory;
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ state: Arc<RwLock<SessionState>>,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<()> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // we may need other infrastructure provided by state, for example:
+ // state.config().get_extension()
Review Comment:
This comment is connected to previous, as getting `config` involves holding read lock. Read may not be as bad as write lock but still it will block writers.
As `create` may take some time (like getting a ML model from repo) read lock may be problematic as well.
Function signature definitely looks better
```rust
async fn create(
&self,
state: &SessionConfig,
statement: CreateFunction,
) -> Result<RegisterFunction>;
```
but lock is still an issue.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1508182996
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -619,6 +624,117 @@ async fn verify_udf_return_type() -> Result<()> {
Ok(())
}
+#[derive(Debug, Default)]
+struct MockFunctionFactory {
+ pub captured_expr: Mutex<Option<Expr>>,
+}
+
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ _config: &SessionConfig,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<RegisterFunction> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // capture expression so we can verify
+ // it has been parsed
+ *self.captured_expr.lock() = statement.params.return_;
+
+ Ok(RegisterFunction::Scalar(Arc::new(mock_udf)))
+ }
+
+ async fn remove(
Review Comment:
will have to follow up on #9304 ... `CREATE FUNCTION` could create UDF or UDaF ... do you maybe have idea how could we decide which unregister_ud*f to call when we get `DROP FUNCTION`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1510213685
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -794,6 +797,55 @@ impl SessionContext {
Ok(false)
}
+ async fn create_function(&self, stmt: CreateFunction) -> Result<DataFrame> {
+ let function = {
+ let state = self.state.read().clone();
+ let function_factory = &state.function_factory;
+
+ match function_factory {
+ Some(f) => f.create(state.config(), stmt).await?,
+ _ => Err(DataFusionError::Configuration(
+ "Function factory has not been configured".into(),
+ ))?,
+ }
+ };
+
+ match function {
+ RegisterFunction::Scalar(f) => {
+ self.state.write().register_udf(f)?;
+ }
+ RegisterFunction::Aggregate(f) => {
+ self.state.write().register_udaf(f)?;
+ }
+ RegisterFunction::Window(f) => {
+ self.state.write().register_udwf(f)?;
+ }
+ RegisterFunction::Table(name, f) => self.register_udtf(&name, f),
+ };
+
+ self.return_empty_dataframe()
+ }
+
+ async fn drop_function(&self, stmt: DropFunction) -> Result<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 {
Review Comment:
That is a good call -- I added a test to sqllogictest
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1509906324
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -1261,7 +1313,38 @@ impl QueryPlanner for DefaultQueryPlanner {
.await
}
}
+/// 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 {
Review Comment:
The last thing standing is trait name, I'm not sure if `FunctionFactory` is the best name for it. It would be appropriate if this PR is part of spark framework :)
any suggestions @alamb ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1509952063
##########
datafusion/expr/src/logical_plan/ddl.rs:
##########
@@ -303,3 +321,66 @@ pub struct DropCatalogSchema {
/// Dummy schema
pub schema: DFSchemaRef,
}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunction {
+ // TODO: There is open question should we expose sqlparser types or redefine them here?
+ // At the moment it make more sense to expose sqlparser types and leave
+ // user to convert them as needed
+ pub or_replace: bool,
+ pub temporary: bool,
+ pub name: String,
+ pub args: Option<Vec<OperateFunctionArg>>,
+ pub return_type: Option<DataType>,
+ // TODO: move this to new struct here
+ pub params: CreateFunctionBody,
+ //pub body: String,
+ /// Dummy schema
+ pub schema: DFSchemaRef,
+}
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct OperateFunctionArg {
+ // it is not really supported so no need to have it here
+ // currently
+ // pub mode: Option<ArgMode>,
+ pub name: Option<Ident>,
+ pub data_type: DataType,
+ pub default_expr: Option<Expr>,
+}
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunctionBody {
+ /// LANGUAGE lang_name
+ pub language: Option<Ident>,
+ /// IMMUTABLE | STABLE | VOLATILE
+ pub behavior: Option<Volatility>,
+ /// AS 'definition'
+ pub as_: Option<DefinitionStatement>,
+ /// RETURN expression
+ pub return_: Option<Expr>,
+}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub enum DefinitionStatement {
+ SingleQuotedDef(String),
+ DoubleDollarDef(String),
+}
+
+impl From<sqlparser::ast::FunctionDefinition> for DefinitionStatement {
+ fn from(value: sqlparser::ast::FunctionDefinition) -> Self {
+ match value {
+ sqlparser::ast::FunctionDefinition::SingleQuotedDef(s) => {
+ Self::SingleQuotedDef(s)
+ }
+ sqlparser::ast::FunctionDefinition::DoubleDollarDef(s) => {
+ Self::DoubleDollarDef(s)
+ }
+ }
+ }
+}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct DropFunction {
+ pub name: String,
+ pub if_exists: bool,
+ pub schema: DFSchemaRef,
Review Comment:
`DdlStatement::schema` needs it. I've tried to remove it and create it when needed but compiler complains it does not live long enough. Can remove `pub`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1976339860
Thanks for all your help @alamb!
Please let me know when new sqlparser lands so I can update dialect used in test
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1508176155
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -619,6 +624,117 @@ async fn verify_udf_return_type() -> Result<()> {
Ok(())
}
+#[derive(Debug, Default)]
+struct MockFunctionFactory {
+ pub captured_expr: Mutex<Option<Expr>>,
+}
+
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ _config: &SessionConfig,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<RegisterFunction> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // capture expression so we can verify
+ // it has been parsed
+ *self.captured_expr.lock() = statement.params.return_;
+
+ Ok(RegisterFunction::Scalar(Arc::new(mock_udf)))
+ }
+
+ async fn remove(
Review Comment:
> lets say we have a SQL UDFs which are expressed by expr, we would need to have some kind of registry (which is not udf_registry) which can be consulted by plan analyser.
Another potential approach would be to encapsulate the code directly into the ScalarUDF and then use an API like `simplify()` that @jayzhan211 and I are working on in https://github.com/apache/arrow-datafusion/pull/9304 to replace the `ScalarUDF` with its `Expr` equivelnt...
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -619,6 +624,117 @@ async fn verify_udf_return_type() -> Result<()> {
Ok(())
}
+#[derive(Debug, Default)]
+struct MockFunctionFactory {
+ pub captured_expr: Mutex<Option<Expr>>,
+}
+
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ _config: &SessionConfig,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<RegisterFunction> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // capture expression so we can verify
+ // it has been parsed
+ *self.captured_expr.lock() = statement.params.return_;
+
+ Ok(RegisterFunction::Scalar(Arc::new(mock_udf)))
+ }
+
+ async fn remove(
Review Comment:
> lets say we have a SQL UDFs which are expressed by expr, we would need to have some kind of registry (which is not udf_registry) which can be consulted by plan analyser.
Another potential approach would be to encapsulate the code directly into the ScalarUDF and then use an API like `simplify()` that @jayzhan211 and I are working on in https://github.com/apache/arrow-datafusion/pull/9304 to replace the `ScalarUDF` with its `Expr` equivalent...
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1509952350
##########
datafusion/expr/src/logical_plan/ddl.rs:
##########
@@ -303,3 +321,66 @@ pub struct DropCatalogSchema {
/// Dummy schema
pub schema: DFSchemaRef,
}
+
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+pub struct CreateFunction {
+ // TODO: There is open question should we expose sqlparser types or redefine them here?
+ // At the moment it make more sense to expose sqlparser types and leave
+ // user to convert them as needed
+ pub or_replace: bool,
+ pub temporary: bool,
+ pub name: String,
+ pub args: Option<Vec<OperateFunctionArg>>,
+ pub return_type: Option<DataType>,
+ // TODO: move this to new struct here
+ pub params: CreateFunctionBody,
+ //pub body: String,
+ /// Dummy schema
+ pub schema: DFSchemaRef,
Review Comment:
If you refer to `schema` same comment like above
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1974786736
and please squash commit on merge, thanks
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1974874013
> Thanks @milenkovicm I would say it is nice to get also tests what will happen if the function is invalid/not compilable or something
>
> another question if function can call another function?
Thanks for your comments @comphead ,
can you please elaborate your comment, thanks
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1510016373
##########
datafusion/core/tests/sql/sql_api.rs:
##########
@@ -16,6 +16,7 @@
// under the License.
use datafusion::prelude::*;
+
use tempfile::TempDir;
Review Comment:
is it related?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1509957293
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -635,6 +640,102 @@ async fn verify_udf_return_type() -> Result<()> {
Ok(())
}
+#[derive(Debug, Default)]
+struct MockFunctionFactory {
+ pub captured_expr: Mutex<Option<Expr>>,
+}
+
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction] statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ _config: &SessionConfig,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<RegisterFunction> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+ datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+ datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) => Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // capture expression so we can verify
+ // it has been parsed
+ *self.captured_expr.lock() = statement.params.return_;
+
+ Ok(RegisterFunction::Scalar(Arc::new(mock_udf)))
+ }
+}
+
+#[tokio::test]
+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();
+ let session_config =
+ SessionConfig::new().set_str("datafusion.sql_parser.dialect", "PostgreSQL");
+
+ let state =
+ SessionState::new_with_config_rt(session_config, Arc::new(runtime_environment))
+ .with_function_factory(function_factory.clone());
+
+ let ctx = SessionContext::new_with_state(state);
+
+ let sql = r#"
+ CREATE FUNCTION better_add(DOUBLE, DOUBLE)
+ RETURNS DOUBLE
+ RETURN $1 + $2
+ "#;
+ let _ = ctx.sql(sql).await?;
Review Comment:
Oh, yes, sorry -- I meant it to be a suggestion for a comment. Sorry about that
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1970719327
Comments totally make sense thanks @alamb,
it would involve a bit of creativity to nail the trait interface, let's try. I'm a bit tied up over next few days will follow up over the weekend probably
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1970860476
Items to follow up (new issues or discussion)
- [ ] Create an example and documents?
- [ ] Should `Generic` dialect support `CREATE STATEMENT`?
- [ ] Add unregister function in `SessionContext`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1506768428
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -1246,7 +1273,42 @@ 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");
+/// ```
+#[async_trait]
+pub trait FunctionFactory: Sync + Send {
+ // TODO: I don't like having RwLock Leaking here, who ever implements it
+ // has to depend ot `parking_lot`. I'f we expose &mut SessionState it
+ // may keep lock of too long.
+ //
+ // Not sure if there is better approach.
+ //
+
+ /// Handles creation of user defined function specified in [CreateFunction] statement
+ async fn create(
+ &self,
+ state: Arc<RwLock<SessionState>>,
Review Comment:
Should we want to implement sql user defined functions, we would need extension to cache logical expressions and provide them to logical analyser so it can replace udf call with the expression.
Id argue we need state or configuration as a parameter as well. Let me do other changes, as they are straightforward and we can try to find most ergonomic function
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1979706229
Thanks @alamb
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add plugable handler for `CREATE FUNCTION` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#issuecomment-1979788522
Thanks again @milenkovicm -- this is a pretty epic feature
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org