Support compute return types from argument values (not just their DataTypes) #8985
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -17,12 +17,13 @@ | |
|
|
||
| //! [`ScalarUDF`]: Scalar User Defined Functions | ||
|
|
||
| use crate::ExprSchemable; | ||
| use crate::{ | ||
| ColumnarValue, Expr, FuncMonotonicity, ReturnTypeFunction, | ||
| ScalarFunctionImplementation, Signature, | ||
| }; | ||
| use arrow::datatypes::DataType; | ||
| use datafusion_common::{ExprSchema, Result}; | ||
| use std::any::Any; | ||
| use std::fmt; | ||
| use std::fmt::Debug; | ||
|
|
@@ -110,7 +111,7 @@ impl ScalarUDF { | |
| /// | ||
| /// If you implement [`ScalarUDFImpl`] directly you should return aliases directly. | ||
| pub fn with_aliases(self, aliases: impl IntoIterator<Item = &'static str>) -> Self { | ||
| Self::new_from_impl(AliasedScalarUDFImpl::new(self.inner.clone(), aliases)) | ||
| } | ||
|
|
||
| /// Returns a [`Expr`] logical expression to call this UDF with specified | ||
|
|
@@ -146,10 +147,17 @@ impl ScalarUDF { | |
| } | ||
|
|
||
| /// The datatype this function returns given the input argument input types. | ||
| /// This function is used when the input arguments are [`Expr`]s. | ||
| /// | ||
| /// | ||
| /// See [`ScalarUDFImpl::return_type_from_exprs`] for more details. | ||
alamb marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| pub fn return_type_from_exprs( | ||
| &self, | ||
| args: &[Expr], | ||
| schema: &dyn ExprSchema, | ||
| ) -> Result<DataType> { | ||
| // If the implementation provides a return_type_from_exprs, use it | ||
| self.inner.return_type_from_exprs(args, schema) | ||
| } | ||
|
|
||
| /// Invoke the function on `args`, returning the appropriate result. | ||
|
|
@@ -249,6 +257,43 @@ pub trait ScalarUDFImpl: Debug + Send + Sync { | |
| /// the arguments | ||
| fn return_type(&self, arg_types: &[DataType]) -> Result<DataType>; | ||
|
|
||
| /// What [`DataType`] will be returned by this function, given the | ||
| /// arguments? | ||
| /// | ||
| /// Note most UDFs should implement [`Self::return_type`] and not this | ||
| /// function. The output type for most functions only depends on the types | ||
| /// of their inputs (e.g. `sqrt(f32)` is always `f32`). | ||
| /// | ||
| /// By default, this function calls [`Self::return_type`] with the | ||
| /// types of each argument. | ||
| /// | ||
| /// This method can be overridden for functions that return different | ||
| /// *types* based on the *values* of their arguments. | ||
| /// | ||
| /// For example, the following two function calls get the same argument | ||
| /// types (something and a `Utf8` string) but return different types based | ||
| /// on the value of the second argument: | ||
| /// | ||
| /// * `arrow_cast(x, 'Int16')` --> `Int16` | ||
| /// * `arrow_cast(x, 'Float32')` --> `Float32` | ||
| /// | ||
| /// # Notes: | ||
| /// | ||
| /// This function must consistently return the same type for the same | ||
| /// logical input even if the input is simplified (e.g. it must return the same | ||
| /// value for `('foo' | 'bar')` as it does for ('foobar'). | ||
|
Comment on lines
+268
to
+292
Contributor
There was a problem hiding this comment. Maybe add some documentation about what would happen if a user tries to implement both And what the suggested implementation for |
||
| fn return_type_from_exprs( | ||
|
Contributor
There was a problem hiding this comment. Since we need to use the we could change the trait impl to something like this pub trait ScalarUDFImpl: Debug + Send + Sync {
/// What [`DataType`] will be returned by this function, given the types of
/// the expr arguments
fn return_type_from_exprs(
&self,
arg_exprs: &[Expr],
schema: &dyn ExprSchema,
) -> Option<Result<DataType>> {
// The default implementation returns None
// so that people don't have to implement `return_type_from_exprs` if they dont want to
None
}
}then change the impl ScalarUDF
/// The datatype this function returns given the input argument input types.
/// This function is used when the input arguments are [`Expr`]s.
/// See [`ScalarUDFImpl::return_type_from_exprs`] for more details.
pub fn return_type_from_exprs<S: ExprSchema>(
&self,
args: &[Expr],
schema: &S,
) -> Result<DataType> {
// If the implementation provides a return_type_from_exprs, use it
if let Some(return_type) = self.inner.return_type_from_exprs(args, schema) {
return_type
// Otherwise, use the return_type function
} else {
let arg_types = args
.iter()
.map(|arg| arg.get_type(schema))
.collect::<Result<Vec<_>>>()?;
self.return_type(&arg_types)
}
}
}this way we don't need to constrain the ScalarFunctionDefinition::UDF(fun) => {
Ok(fun.return_type_from_exprs(&args, schema)?)
}and it still makes It does make it very slightly less ergonomic as end users now need to wrap their body in an
Contributor
Author
There was a problem hiding this comment. This works well on my side. Thanks! Another question for me is, how can user implement
Contributor
There was a problem hiding this comment. Hmm yeah @yyy1000 you still run into the same error then. I'm wondering if it'd be easiest to just change the type signature on pub trait ExprSchemable<S: ExprSchema> {
/// given a schema, return the type of the expr
fn get_type(&self, schema: &S) -> Result<DataType>;
/// given a schema, return the nullability of the expr
fn nullable(&self, input_schema: &S) -> Result<bool>;
/// given a schema, return the expr's optional metadata
fn metadata(&self, schema: &S) -> Result<HashMap<String, String>>;
/// convert to a field with respect to a schema
fn to_field(&self, input_schema: &DFSchema) -> Result<DFField>;
/// cast to a type with respect to a schema
fn cast_to(self, cast_to_type: &DataType, schema: &S) -> Result<Expr>;
}
impl ExprSchemable<DFSchema> for Expr {
//...
}then the trait can just go back to the original implementation you had using fn return_type_from_exprs(
&self,
arg_exprs: &[Expr],
schema: &DFSchema,
) -> Result<DataType> {
let arg_types = arg_exprs
.iter()
.map(|e| e.get_type(schema))
.collect::<Result<Vec<_>>>()?;
self.return_type(&arg_types)
}I tried this locally and was able to get things to compile locally, and was able to implement a udf using the trait. It does make it a little less flexible as it's expecting a I think the only other approach would be to make
Contributor
Author
There was a problem hiding this comment. Thanks for your help! @universalmind303
Contributor
There was a problem hiding this comment. Update is I changed the signature to take |
||
| &self, | ||
| args: &[Expr], | ||
| schema: &dyn ExprSchema, | ||
| ) -> Result<DataType> { | ||
| let arg_types = args | ||
| .iter() | ||
| .map(|arg| arg.get_type(schema)) | ||
| .collect::<Result<Vec<_>>>()?; | ||
| self.return_type(&arg_types) | ||
| } | ||
|
|
||
| /// Invoke the function on `args`, returning the appropriate result | ||
| /// | ||
| /// The function will be invoked passed with the slice of [`ColumnarValue`] | ||
|
|
@@ -290,13 +335,13 @@ pub trait ScalarUDFImpl: Debug + Send + Sync { | |
| /// implement [`ScalarUDFImpl`], which supports aliases, directly if possible. | ||
| #[derive(Debug)] | ||
| struct AliasedScalarUDFImpl { | ||
| inner: Arc<dyn ScalarUDFImpl>, | ||
| aliases: Vec<String>, | ||
| } | ||
|
|
||
| impl AliasedScalarUDFImpl { | ||
| pub fn new( | ||
| inner: Arc<dyn ScalarUDFImpl>, | ||
| new_aliases: impl IntoIterator<Item = &'static str>, | ||
| ) -> Self { | ||
| let mut aliases = inner.aliases().to_vec(); | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Here is an example of the feature working