Support `Expr` creation for `ScalarUDF`: Resolve function calls by name during planning

[2010YOUY01]

Is your feature request related to a problem or challenge?

Motivation

There is ongoing work migrating BuitlinScalarFunctions -> ScalarUDF #8045, and we noticed one Expr related issue during prototyping:

We can use Expr API to create builtin scalar functions directly:

let expr1 = abs(lit(-1));
let expr2 = call_fn("abs", vec![lit(-1)]);

We want to still use this API after functions are migrated to ScalarUDF based implementation, it's not possible now because those Expr creations are stateless, and ScalarUDFs are registered inside SessionState. It's only possible if we do

let ctx = create_context()?;
ctx.register_udf(abs_impl);
let expr2 = ctx.call_fn("abs", vec![lit(-1)]);

Describe the solution you'd like

To continue supporting the original stateless Expr creation API, we can create Expr with only the string function name, and resolve the function during logical optimization.
Then call_fn("abs", vec![lit(-1)]) can be supported (now it only support BuitlinScalarFunction and don't support UDFs)
Another macro-based API abs(lit(-1)) can be supported if we hard code all possible function names within the core (should we do that?)

Potential implementation is:

1. After let expr2 = call_fn("abs", vec![lit(-1)]);, create a ScalarUDF expression with dummy implementation.
2. Add an AnalyzerRule(a mandatory logical optimizer rule) to resolve this UDF name using external functions registered in SessionState

Issue:

Now function implementation is inside SessionState but outside SessionConfig, and the analyzer can only access SessionConfig.
We have to move the function registry into SessionConfig first (or is there any better way?)

pub struct SessionState {
    // ...
    /// Scalar functions that are registered with the context
    scalar_functions: HashMap<String, Arc<ScalarUDF>>,
    /// Session configuration
    config: SessionConfig,

// ...
// `self.options()` is `ConfigOption`
            let analyzed_plan =
                self.analyzer
                    .execute_and_check(plan, self.options(), |_, _| {})?;

cc @alamb I wonder if you have any thoughts on this approach 🤔

Describe alternatives you've considered

No response

Additional context

No response

[alamb]

Thanks @2010YOUY01 !

We have to move the function registry into SessionConfig first (or is there any better way?)

What if we changed how Expr::ScalarFunction looks (and remove Expr::ScalarUDF:

enum Expr {
...
    /// Represents the call of a built-in, or UDF scalar function with a set of arguments.
    ScalarFunction(ScalarFunction),
...
}

Instead of

#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct ScalarFunction {
    /// The function
    pub fun: built_in_function::BuiltinScalarFunction,
    /// List of expressions to feed to the functions as arguments
    pub args: Vec<Expr>,
}

Make it look like

pub enum ScalarFunctionDefinition {
  /// Resolved to a built in scalar function
  /// (will be removed long term)
  BuiltIn(built_in_function::BuiltinScalarFunction),
  /// Resolved to a user defined function
  UDF(ScalarUDF),
  /// A scalar function that will be called by name
  Name(Arc<str>),
}

#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct ScalarFunction {
    /// The function
    pub fun: ScalarFunctionDefinition,
    /// List of expressions to feed to the functions as arguments
    pub args: Vec<Expr>,
}

And that way an expr function like abs could look like

fn abs(arg: Expr) -> Expr {
  Expr::ScaarFunction {
    fun: ScalarFunctionDefintion::Name("abs".into()),
    args: vec![arg],
  }
}

I am not sure how large of a change this would be -- we would have to try it and see what it looked like.

[alamb]

, and the analyzer can only access SessionConfig.

I think passing a impl FunctionRegistry to the analyzer rule would be another way (so we could pass SessionState to that rule without it having to know about the actual SessionState

[alamb]

@2010YOUY01 I wonder if you still plan to work on this item? If not, no worries I can do it, but I wanted to check with you before doing so

[2010YOUY01]

@2010YOUY01 I wonder if you still plan to work on this item? If not, no worries I can do it, but I wanted to check with you before doing so

Please feel free to proceed if you would like it to get done sooner, thank you @alamb 🙏🏼
I will catch up next week when I have more available time to help

[edmondop]

@alamb I checked out and on main the structs and enums are already written as you recommended

pub enum ScalarFunctionDefinition {
  /// Resolved to a built in scalar function
  /// (will be removed long term)
  BuiltIn(built_in_function::BuiltinScalarFunction),
  /// Resolved to a user defined function
  UDF(ScalarUDF),
  /// A scalar function that will be called by name
  Name(Arc<str>),
}

#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct ScalarFunction {
    /// The function
    pub fun: ScalarFunctionDefinition,
    /// List of expressions to feed to the functions as arguments
    pub args: Vec<Expr>,
}

So I am not sure I understood correctly. I could see that call_fn resolves at invocation site rather than at planning site, and changing the behavior would probably mean changing this code. What is the right type of Expr this code should return?

/// Calls a named built in function
/// ```
/// use datafusion_expr::{col, lit, call_fn};
///
/// // create the expression sin(x) < 0.2
/// let expr = call_fn("sin", vec![col("x")]).unwrap().lt(lit(0.2));
/// ```
pub fn call_fn(name: impl AsRef<str>, args: Vec<Expr>) -> Result<Expr> {
    match name.as_ref().parse::<BuiltinScalarFunction>() {
        Ok(fun) => Ok(Expr::ScalarFunction(ScalarFunction::new(fun, args))),
        Err(e) => Err(e),
    }
}

[alamb]

So I am not sure I understood correctly. I could see that call_fn resolves at invocation site rather than at planning site, and changing the behavior would probably mean changing this code. What is the right type of Expr this code should return?

I think the idea would be to make it be something like this

/// Calls a named function
/// ```
/// use datafusion_expr::{col, lit, call_fn};
///
/// // create the expression sin(x) < 0.2
/// let expr = call_fn("sin", vec![col("x")]).unwrap().lt(lit(0.2));
/// ```
pub fn call_fn(name: impl AsRef<str>, args: Vec<Expr>) -> Result<Expr> {
     Ok(Expr::ScalarFunction(ScalarFunctionDefinition::Name(Arc::new(name.as_ref()))
}```

And then add code to replace instances of `ScalarFunctionDefinition::Name` with either `ScalarFunctionDefinition::BuiltIn` or `ScalarFunction::Udf` to the optimizer. 

Perhaps you can use the same logic here: https://github.com/apache/arrow-datafusion/blob/d9d8ddd5f770817f325190c4c0cc02436e7777e6/datafusion/sql/src/expr/function.rs#L66-L76

BTW the point of doing this is so we can remove `BuiltInScalarFunction` eventually  incrementally and treat all functions the same