Higher-order functions and async UDF don't work together

[gstvg]

Describe the bug

Async udfs are removed from their original position and evaluated within a dedicated execution plan, where lambda variables children of the async udf aren't available. For example, select array_transform([-2], v -> async_abs(v)) is planned into:

logical_plan
01)Projection: array_transform(List([-2]), (v) -> async_abs(v)) AS array_transform(make_array(Int64(-2)),(v) -> async_abs(v))
02)--EmptyRelation: rows=1
physical_plan
01)ProjectionExec: expr=[array_transform([-2], (v) -> __async_fn_0@0) as array_transform(make_array(Int64(-2)),(v) -> async_abs(v))]
02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
03)----AsyncFuncExec: async_expr=[async_expr(name=__async_fn_0, expr=async_abs(v@0))]
04)------PlaceholderRowExec

The async_abs(v@0) in AsyncFuncExpr contains a lambda variable v@0 which isn't available there

Also, using lambda variables and async udf without lambda variables don't work either: select array_transform([1], v -> v + async_abs(-2))

logical_plan
01)Projection: array_transform(List([1]), (v) -> v + async_abs(Int64(-2))) AS array_transform(make_array(Int64(1)),(v) -> v + async_abs(Int64(-2)))
02)--EmptyRelation: rows=1
physical_plan
01)ProjectionExec: expr=[array_transform([1], (v) -> v@0 + __async_fn_0@0) as array_transform(make_array(Int64(1)),(v) -> v + async_abs(Int64(-2)))]
02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
03)----AsyncFuncExec: async_expr=[async_expr(name=__async_fn_0, expr=async_abs(-2))]
04)------PlaceholderRowExec

Because the new __async_fn_0 column isn't present in the schema during physical planning, lambda variable get the same index: v@0 + __async_fn_0@0, and generates an error during execution (I believe this is similar to #18149)

To Reproduce

SELECT array_transform([1], v -> async_udf(v)) or select array_transform([1], v -> v + async_abs(-2))

Expected behavior

Async udfs should work normally regardless of using lambda variables or not

Additional context

Async UDFs PR #14837
Higher-order function PR #21679

[gstvg]

I can imagine 3 solutions:

• Make HigherOrderUDF::evaluate and LambdaArgument::evaluate async, so that any async_udf within a lambda body can be evaluated before the rest of the body, it's output added to the lambda batch, the async_udf rewrited into a column, and then the whole body evaluated. All higher-order functions would be moved to AsyncFuncExec as well.
• Add a default, blocking execution for async udf that uses lambda variables using tokio::task::block_in_place, like below, but it won't work on WASM:

    fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
        block_in_place(move || {
            Handle::try_current()
                .map_err(|err| {
                    exec_datafusion_err!(
                        "{} can't get handle of current tokio runtime: {err}",
                        self.name()
                    )
                })?
                .block_on(async {
                    invoke_with_args(
                        self,
                        Arc::clone(&self.return_field),
                        batch,
                        Arc::clone(&self.config_options),
                    )
                    .await
                })
        })
    }

• Swap HigherOrderUDF::evaluate for the following method, so that AsyncFuncExec can handle any async udf with a lambda body for each of it's invocation. Higher-order functions with async_udf within it's body would be moved to AsyncFuncExec

trait HigherOrderUDF {
    fn make_evaluator(&self, args: HigherOrderFunctionArgs) -> Box<dyn Evaluator>;
}

enum Step {
    /// AsyncFuncExec will evaluate the async nodes of the expr using the batch, append their outputs to batch and
    /// rewrite the async nodes to column ref,then evaluate the whole expr with it,
    /// and provide the result to the next Evaluator::next call
    Partial((RecordBatch, Arc<dyn PhysicalExpr>)),
    Complete(ColumnarValue)
}

trait Evaluator {
    fn next(&mut self, last_step_output: Option<ColumnarValue>) -> Result<Step>;
}

I'll explore the first option, but I'm open to suggestions
cc @rluvaton @LiaCastaneda @comphead @pepijnve @martin-g

[rluvaton]

I'm against the first option as I don't wanna add overhead of async for not very common case of async Udf let alone async Udf inside lambda

[rluvaton]

We can start by first disallowing async functions in lambda until we solve this

[gstvg]

I'm against the first option as I don't wanna add overhead of async for not very common case of async Udf let alone async Udf inside lambda

Thats true. I don't like any of those options actually, the second is blocking and don't work on WASM, and the last option make the API strange. PR disallowing async functions in lambda opened at #22097