UDF: Can I get a return type from a constant parameter in args?

[JasonLi-cn]

Def one UDF:

my_udf(expr, 'return_type')

Using:

case1:

my_duf(col, 'UInt32')

Return Type is:
DataType::UInt32

case2:

my_duf(col, 'UInt64')

Return Type is:
DataType::UInt64

At present, the above udf cannot be realized through the ReturnTypeFunction. This is because the ReturnTypeFunction provides only the data types of the input parameters.

[alamb]

Hi @JasonLi-cn -- I am not sure you can do this as a UDF -- the only function I know that is similar is arrow_cast but that is handled specially in the parser for exactly the reason you describe:

https://github.com/apache/arrow-datafusion/blob/85f3578f5fb47d28a8bc3a7b9be0284b3ced0fcd/datafusion/sql/src/expr/arrow_cast.rs#L31-L73

[JasonLi-cn]

The usage is indeed similar to arrow_cast, but the functionality is different. If I want to implement this function, can I only change the datafusion source code?

[alamb]

The usage is indeed similar to arrow_cast, but the functionality is different. If I want to implement this function, can I only change the datafusion source code?

The other alternative is to figure out a way to extend the UDF API to support the interface https://docs.rs/datafusion/latest/datafusion/physical_expr/udf/struct.ScalarUDF.html#structfield.return_type

Maybe we could use some trait trickery to change from

return_type: Arc<dyn Fn(&[DataType]) -> Result<Arc<DataType, Global>, DataFusionError> + Sync + Send, Global>

to

return_type: Arc<dyn ReturnType>

with

trait ReturnType {
  fn from_datatypes(types: &[DataType]) -> Result<Arc<DataType>, DataFusionError>);

  // calculate return type from the actual arguments (provides a way to handle constants)
  fn from_args(args: &[Expr]) -> Result<Arc<DataType>, DataFusionError>);

}

/// implement the trait for all existing functions
impl ReturnType for dyn Fn(&[DataType]) -> Result<Arc<DataType, Global>, DataFusionError> + Sync + Send {
...
}

[JasonLi-cn]

That's a good idea. But in addition, we need to change code where we call return_type:

https://github.com/apache/arrow-datafusion/blob/32fe1761287c9b58294a54805a0fcafc4ab046b8/datafusion/expr/src/expr_schema.rs#L84-L89

and we need to set the priorities of from_datatypes and from_args.

[JasonLi-cn]

In addition, args also includes PhysicalExpr.

https://github.com/apache/arrow-datafusion/blob/a64f36d8fbce35cb7146271fb5d9ab98a9f094e6/datafusion/physical-expr/src/udf.rs#L27-L43

[JasonLi-cn]

I have done a simple test, please help review the code. @alamb

main...JasonLi-cn:arrow-datafusion:feature/refactor_udf

[alamb]

I thinnk the idea of

/// Factory that returns the functions's return type given the input argument types and constant arguments
pub trait ReturnTypeFactory: Send + Sync {
    fn infer(
        &self,
        data_types: &[DataType],
        literals: &[(usize, ScalarValue)],
    ) -> Result<Arc<DataType>>;
}


impl ReturnTypeFactory for ReturnTypeFunction {
    fn infer(
        &self,
        data_types: &[DataType],
        _literals: &[(usize, ScalarValue)],
    ) -> Result<Arc<DataType>> {
        self(data_types)
    }
}

And then

/// Logical representation of a UDF.
#[derive(Clone)]
pub struct ScalarUDF {
...
    /// Return type
    pub return_type: Arc<dyn ReturnTypeFactory>,
...
}

Makes sense to me @JasonLi-cn and is a nice design

It looks to me like it would be backwards compatible (in the sense that all existing ScalarUDFs would work without modification). Is that your understanding as well (the CI tests would definitely cover this).

Also cc @2010YOUY01 who is looking into other changes in the ScalarUDF interfaces: #7752

[alamb]

I filed #8624 to track this request and so I am going to mark this discussion closed