Add support for lambda/higher order functions

[gstvg]

Is your feature request related to a problem or challenge?

Some engines, like DuckDB and Clickhouse, supports lambda functions, like:

SELECT list_filter(numbers, x -> x > 40) as greater_than_40s FROM relation

There's already support for the syntax in sqlparser-rs

From #12206 (comment) by @jayzhan211

Describe the solution you'd like

One of

1. Add Expr::Lambda{arg_names: Vec<String>, expr: Expr} variant
2. Change ScalarFunction.args :

struct ScalarFunction {
    args: Vec<ScalarFunctionArgument>
    ... omitted
}

enum ScalarFunctionArgument {
    Expr(Expr),
    Lambda{ arg_names: Vec<String>, body: Expr }
}

Create a LambdaPhysicalExpr that holds the lambda physical expr and returns ScalarValue::Null in PhysicalExpr::evaluate or an error(similar to NoOp and UnkownColumn)
Add ScalarUDFImpl::lambdas_schema to return the types of the args of all lambda arguments it receives, so a input schema can be built and used to generated the LambdaPhysicalExpr

And then, one of:

1. Make ScalarFunctionArgs generic over the arg type, with ColumnarValue as default, and add a new ScalarUDF method invoke_with_lambda_args / invoke_higher_order_with_args that receives ScalarFunctionArgs<ColumnarValueOrLambda> instead of ScalarFunctionArgs, and has a default implementation calling invoke_with_args, returning an error if any arg is an lambda. The lambda arg is created if any children PhysicalExpr is a LambdaPhysicalExpr

enum ColumnarValueOrLambda {
    Value(ColumnarValue),
    Lambda(&dyn PhysicalExpr)
}

struct ScalarFunctionArgs<T = ColumnarValue> {
    args: Vec<T>,
    ... omitted
}

2. Add physical_exprs: Vec<Arc<dyn PhysicalExpr>> to ScalarFunctionArgs, so lambda expressions can be extracted from
3. Add LambdaScalarUDF trait (Duplication with ScalarUDF, more work and more code to mantain, less flexible than 1, one more trait to document and to users to reason about)
4. Change ScalarFunctionArgs.args to Vec<ColumnarValueOrLambda> instead of Vec<ColumnarValue> (A lot of breakage, including public)
5. Add Lambda variant to ColumnarValue (Even more breakge than 2, and a lambda doesn't quite fit the concept of a columnar value)

Describe alternatives you've considered

No response

Additional context

No response

[rkrishn7]

Hello! I'm interested in working on this.

Agreed that tweaking the existing ScalarUDF code would be ideal as it would be least breaking. However, after poking around a bit, I have some concerns:

• How do implementers define a signature for a UDF that takes one or more lambdas? Do we just skip type checking of lambda arguments? How does this affect the ordering?
• How should one map lambda function arguments to concrete values in the body? If the function only operates on a list, this becomes easier to express.

A more prohibitive interface certainly seems like it would help implementation here, albeit at the expense of consumer flexibility.

Any thoughts/feedback would be much appreciated. Thanks!

[gstvg]

Hi @rkrishn7, have you already started any work?

Since creating this issue, I’ve been working on a PoC to refine the proposal, and I managed to get it working. I’d like to continue with it if you’re OK with that—I probably should have assigned the issue to myself earlier, my bad 🤦🏻‍♂️! Anyway, many thanks for your interest and for your questionings, and I would really appreciate your input in futures questions (I expect many to arise!)

How do implementers define a signature for a UDF that takes one or more lambdas? Do we just skip type checking of lambda arguments? How does this affect the ordering?

I added a new ScalarUDF method invoke_with_lambda_args that receives ScalarFunctionArgs instead of ScalarFunctionArgs, and has a default implementation calling invoke_with_args, returning an error if any arg is a lambda
This should preserve ordering of any arg.
But maybe we could add an lambda_arguments to ScalarFunctionArgs similtar to here, what do you think ?

How should one map lambda function arguments to concrete values in the body? If the function only operates on a list, this becomes easier to express.

The lambda implementation should build a RecordBatch, likely using some other concrete argument, and pass it to the lambda PhysicalExpr evaluate method

That said, if you have any thoughts on the current interface, I’d love to hear your feedback. It's a list_map([1,2], x -> x*2) udf, you can see the executed output here. Thanks! 🙏

[gstvg]

Lambda syntax only works with databricks dialect, proposed fix here apache/datafusion-sqlparser-rs#1686

[rkrishn7]

Hi @gstvg! No worries at all, please continue working on it since you've started!

Happy to take a look at the work you've done already and give feedback where it makes sense 👍🏾

[gstvg]

take

[alexanderbianchi]

@gstvg Is there an update on the status here?

[gstvg]

Hi @alexanderbianchi. The code is already functional, mostly finished, but it ended up being a bigger change than I expected, ~3k LOC, some of it in core parts of DF. After I got stuck documenting and splitting it into smaller PRs, I stopped for a while.. but it's already been 6 months (!!). I'll start working on it again this weekend and hopefully we can make it into the 51.0 release 🙏