Thanks a lot @gstvg, huge amount of work apparently (even if I am not enough skilled to judge). I am waiting for such functions since a very long time.. IMHO, this is a MUST HAVE in datafusion. All serious analytics engine provide this kind of functions, DuckDB implementation is very mature, Polars seems to have it as well even if I didn"t tried it yet and I really hope that datafusion will jump into this enhancement very quickly as it is again a big miss.

Use cases like list_filter(list, expr), list_map(list, expr) and list_reduce(list, base, expr) will be available soon.

Again, thanks a lot for your work and I cross the fingers for a quick progress.

@gstvg This is super exciting! I'll review over the next couple of days.

cc @SparkApplicationMaster @andygrove would be great to get your thoughts too 😃

Thanks @fbx31. This is high priority for me right now. I also hope we can finish this quickly. It's indeed a lot work, but now I believe it ended up that way because I went to fast in the wrong direction... More on the comment below. Thanks again!

Thanks @shehabgamin. I analyzed the Spark implementation one last time, and realized I could have done similarly since the beginning... I've already tested it locally and pretend to push soon, the diff reduced to ~2K LOC, mostly in new, self contained code with small changes in existing code and without requiring changes in downstream code. I hope you haven't put time reviewing it yet because I believe the spark approach is much better and also easier/faster to review

Basically, for planning, we lazily set the Field in the lambda column itself, instead of injecting it into a DFSchema

//logical
struct LambdaColumn {
    name: String,
    field: Option<FieldRef>,
}

And for evaluation, again, we lazily set the value of a lambda column instead of injecting it into a RecordBatch

//physical
struct LambdaColumn {
    name: String,
    field: FieldRef,
    value: Option<ColumnarValue>,
}

I think I initially ruled that out because other expressions doesn't contain it's own Field, and this have been discussed before in #12604 *, and didn't went forward, which the difference that the other expressions Field's naturally exist in the plan schema, which now I believe justifies this difference.
And most importantly, because I had no idea of how much work and downstream churn injecting the lambdas parameters into the Schema/RecordBatch would cause 🤦

I will push soon and update the PR description, thanks again!

*This is being discussed again in #18845

This is ready for a first look @shehabgamin 🚀

Hello @gstvg, thanks for your last submission, some questions came to my mind reading your proposal, I see quite well the map/filter operations that are producing array in result but I am a bit "confused" about the support of "reduce" like operation that could produce any type as output like another array or scalar values with an aggegation value that has to be first intialized by an Expr through the call to HoF and the "propagation" of this value accross the plan architecture of datafusion.
I think about use cases like array_reduce(input: [], init_aggregate_val: Expr, lamba: Expr) with lambda(current_value_of_input, aggregate_value) -> aggregate_value
Is it possible with your proposal ? (I saw potential several parameters for the lambda so I suppose your proposal is compliant ?)
Also, is it possible to "nest" the HoF like: array_transform(arr_transform([], lamda1), lamba2) ?
For this last question, I don't see any "obstacles" but I prefer to confirm with you... :-)
And also, pratically, how to reference "current_value" and "aggregate_value" in the dataframe API? These are not named in the schema and these are not query parameter as well so it's not clear for me how to reference these values inside the expressions:
df.with_column("new_col", array_transform(col("my_array"), max( * 2, 1000)))
Sorry to question like that, but I was bumping my head on the walls on this topic since months without success (not enough mastering datafusion logical/physical plans logic).

Many, many thanks !!!

Hi @fbx31, reduce is supported too. It isn't included here because it's implementation is not as simple as array_transform and would make this PR even bigger. But if any reviewer/maintainer request it we can add it and make sure it works since the beginning.

Nesting is supported, some tests already cover it

With the new approach you need to use lambda_col("current_value") instead of col("current_value"):

df.with_column(
    "new_col",
    array_transform(
        col("my_array"),
        lambda(
            vec!["current_value"], // define the parameters name 
            lambda_col("current_value") * 2 // the lambda body
        )
    )
)

SQL equivalent:

 array_transform(my_array, current_value -> current_value * 2)

No worries, feel free to ask any questions you have

This is ready for a first look @shehabgamin 🚀

@gstvg Exciting! I've started reviewing, it will take me a few days to get through it all. Thank you for thoroughly explaining all of the options btw, it's super helpful!

Crystal clear @gstvg thanks a lot, cannot wait to test and use it.
A bit disappointed for reduce operation but i fully understand the drawbacks for the reviewers.
Hope this review will be successfull.
And thanks in advance to all reviewers and of course to the author of this PR.

This is an exciting initiative!

I skimmed through the PR, and I have some thoughts regarding how this can fit into DataFusion's existing setup. I wonder if we can have Expr::LambdaFunction(LambdaFunction), similar to Expr::ScalarFunction(ScalarFunction)? Here are my reasoning:

1. LambdaFunction can represent a "resolved" lambda function call in the logical plan. In contrast, Expr::Lambda and Expr::LambdaColumn are only fragments whose data types etc. are not well-defined which may be challenging to work with in other parts of the library (e.g. ExprSchemable).
2. I feel ScalarUDF may not fit cleanly for lambda functions, so we can then have a separate abstraction (e.g. LambdaUDF) inside Expr::LambdaFunction. The ArrayTransform example shows that we would have to use ScalarFunctionArgs which was not originally designed for the lambda use case. (IMO when the function is actually "invoked" with Arrow arrays, the lambda parameter should have been resolved and removed from the argument list.)
3. Once we have a self-contained logical expression, the entire query analysis flow might become easier to reason about (from SqlToRel to Expr::LambdaFunction to a dedicated PhysicalExpr and the corresponding physical planner).

This is just my rough thoughts. Happy to discuss!

Thanks @linhr
I think a new LambdaUDF is reasonable, and I'm not strongly inclined towards using ScalarUDF either

Expr::Lambda and Expr::LambdaColumn are only fragments whose data types etc. are not well-defined which may be challenging to work with in other parts of the library (e.g. ExprSchemable).

Yeah, this is the major challenge in this PR. Currently, Expr::Lambda always return DataType::Null and Expr::LambdaVariable embeds a FieldRef which is used to implement the ExprSchemable and don't even look at the DFSchema

when the function is actually "invoked" with Arrow arrays, the lambda parameter should have been resolved and removed from the argument list.

Could you expand this further? It is like LambdaUDF having a resolve_lambdas method which result is passed to the invoke_with_args method?

Could you expand this further? It is like LambdaUDF having a resolve_lambdas method which result is passed to the invoke_with_args method?

Suppose we have array_transform([1, 2], v -> v*2). We could have a trait LambdaUDF and have impl LambdaUDF for ArrayTransform. (Or we follow the existing convention to have struct LambdaUDF and trait LambdaUDFImpl separately.) A logical representation of v -> v*2 is passed to ArrayTransform::new(). For Expr::LambdaFunction(LambdaFunction), we can have LambdaFunction { func: Arc<dyn LambdaUDF>, args: Vec<Expr> } where the non-lambda parameter [1, 2] is stored in args.

During physical planning, we could resolve ArrayTransform into PhysicalArrayTransform which stores v -> v*2 resolved as certain PhysicalExpr. We have a trait PhysicalLambdaUDF and impl PhysicalLambdaUDF for PhysicalArrayTransform. The trait method PhysicalLambdaUDF::invoke_with_args accepts the Arrow array [1, 2] and compute the results. I'd imagine this invocation can be done in a general (physical) LambdaFunctionExpr that works for all lambda functions, similar to ScalarFunctionExpr.

When I worked with ScalarUDF, I notice that the logic required for logical representation, physical planning, and the actual execution are all within a single ScalarUDFImpl trait. If we look at how these trait methods are used by various planning/execution stages, we might get the big picture how a parallel code structure (with multiple traits to separate the responsibilities) can be designed for lambda functions.

I haven't thought about function registry, documentation etc. which we can get for free in the existing ScalarUDF setup. So some more investigation is needed to estimate the amount of work if we explore the route I described above.

@linhr Thanks, I got it now. I think it's possible to resolve lambda parameters and remove them from the arguments list with ScalarUDF itself, with a few cons compared to a dedicated to LambdaUDF, but with much less code changes and similar enough to compare the alternatives. I believe I'm already finishing it, and will push to another branch soon.

@gstvg any updates on this one? 😇

@linhr
Really really sorry for the delay

I think it's possible to resolve lambda parameters and remove them from the arguments list with ScalarUDF itself, with a few cons compared to a dedicated to LambdaUDF, but with much less code changes and similar enough to compare the alternatives. I believe I'm already finishing it, and will push to another branch soon.

Unfortunately, trying to resolve lambdas with ScalarUDF didn't work well.
So, before moving to a LambdaUDF based implementation, I would like to further discuss the current approach.

Expr::Lambda and Expr::LambdaColumn are only fragments whose data types etc. are not well-defined which may be challenging to work with in other parts of the library (e.g. ExprSchemable).

Yeah, this is the major challenge in this PR. Currently, Expr::Lambda always return DataType::Null and Expr::LambdaVariable embeds a FieldRef which is used to implement the ExprSchemable and don't even look at the DFSchema

I mean, this was the major challenge of the PR on it's first implementation, now that we use LambdaVariable with a Field, I consider this to be solved. And currently Expr::Lambda actually returns the return_field of it's body, but I believe we can also return a field with DataType::Null if deemed better.

LambdaFunction can represent a "resolved" lambda function call in the logical plan. In contrast, Expr::Lambda and Expr::LambdaColumn are only fragments whose data types etc. are not well-defined which may be challenging to work with in other parts of the library (e.g. ExprSchemable).
... when the function is actually "invoked" with Arrow arrays, the lambda parameter should have been resolved and removed from the argument list

Sorry, at first I thought this only means resolving/embedding the lambda body within the UDF itself and removing it from Expr::ScalarFunction.args/LambdaFunction.args, and also removing Expr::Lambda variant from the logical Expr enum and it's physical counterpart, but now I'm not sure it also means omitting the lambda body from TreeNode traversals?

To streamline the discussion, I have a few comments for each of these options, if they're applicable:

Omitting the lambda body from TreeNode traversals

At least the following traversals currently use TreeNode and would require adjust if the body is omitted:

Projection pushdown for column capture, type coercion for correctness, expr simplifier and CSE for performance.

Being only 4 internal traversals, is easy to specially handle them, but since DF is an open system, I believe there should be a way to downstream users to visit/transform lambdas expressions. We could document that they need to be specially handled, outside TreeNode, or we could offer an API for it, and after all, I believe the ideal API would be exactly the TreeNode API.

In my first iteration of this PR, the "outdated" sections of the description and of the first comment, I manually checked every TreeNode usage on logical and physical expressions, and find none where Expr::Lambda should be specially handled(it is simply ignored), and Expr::LambdaVariable only required handling in few places.

Removing Expr::Lambda variant from Expr enum and it's physical counterpart

Removing Expr::Lambda from AST makes some tree traversals more difficult, while keeping it doesn't make difference: it's simply ignored as most expressions in most traversals.

One example is BindLambdaVariable in this PR. Is not important to understand what it does, just that it's bigger, with more boilerplate and harder to read without Expr::Lambda:

struct BindLambdaVariable<'a> {
    variables: HashMap<&'a str, (ArrayRef, usize)>, 
    //(variable value, number of times it has been shadowed by other lambdas variables in inner scopes in the current position of the tree)
}

impl TreeNodeRewriter for BindLambdaVariable<'_> {
    type Node = Arc<dyn PhysicalExpr>;

    fn f_down(&mut self, node: Self::Node) -> Result<Transformed<Self::Node>> {
        if let Some(lambda_variable) = node.as_any().downcast_ref::<LambdaVariable>() {
            if let Some((value, shadows)) = self.variables.get(lambda_variable.name()) {
                if *shadows == 0 {
                    return Ok(Transformed::yes(Arc::new(
                        lambda_variable.clone().with_value(value.clone()),
                    )));
                }
            }
        } else if let Some(inner_lambda) = node.as_any().downcast_ref::<LambdaExpr>() {
            for param in inner_lambda.params() {
                if let Some((_value, shadows)) = self.variables.get_mut(param.as_str()) {
                    *shadows += 1;
                }
            }
        }

        Ok(Transformed::no(node))
    }

    fn f_up(&mut self, node: Self::Node) -> Result<Transformed<Self::Node>> {
        if let Some(inner_lambda) = node.as_any().downcast_ref::<LambdaExpr>() {
            for param in inner_lambda.params() {
                if let Some((_value, shadows)) = self.variables.get_mut(param.as_str()) {
                    *shadows -= 1;
                }
            }
        }

        Ok(Transformed::no(node))
    }
}

Implementation without Physical Lambda on AST

fn bind_lambda_variables(
    node: Arc<dyn PhysicalExpr>,
    params: &mut HashMap<&str, (ArrayRef, usize)>,
) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
    node.transform_down(|node| {
        if let Some(lambda_variable) = node.as_any().downcast_ref::<LambdaVariable>() {
            if let Some((value, shadows)) = params.get(lambda_variable.name()) {
                if *shadows == 0 {
                    return Ok(Transformed::yes(Arc::new(
                        lambda_variable.clone().with_value(value.clone()),
                    )));
                }
            }
        } else if let Some(fun) = node.as_any().downcast_ref::<LambdaUDFExpr>() {
            let mut transformed = false;

            let new_lambdas = fun
                .lambdas()
                .iter()
                .map(|(params_names, body)| {
                    for param in *params_names {
                        if let Some((_value, shadows)) = params.get_mut(param.as_str()) {
                            *shadows += 1;
                        }
                    }

                    let new_body = bind_lambda_variables(Arc::clone(body), params)?;
                    transformed |= new_body.transformed;

                    for param in *params_names {
                        if let Some((_value, shadows)) = params.get_mut(param.as_str()) {
                            *shadows -= 1;
                        }
                    }

                    Ok((params_names.to_vec(), new_body.data))
                })
                .collect::<Result<_>>()?;

            let new_args = fun
                .args()
                .iter()
                .map(|arg| {
                    let new_arg = bind_lambda_variables(Arc::clone(arg), params)?;
                    transformed |= new_arg.transformed;
                    Ok(new_arg.data)
                })
                .collect::<Result<_>>()?;

            let fun = fun.with_new_children(new_args, new_lambdas);

            return Ok(Transformed::new(
                Arc::new(fun),
                transformed,
                TreeNodeRecursion::Stop,
            ));
        }

        Ok(Transformed::no(node))
    })
}

Note that this transformation always return a constant TreeNodeRecursion value. If not, TreeNodeRecursion would also needed to be manually handled, adding more boilerplate that TreeNode handles automatically when Expr::Lambda is present

Resolve/embed lambdas on the UDF implementation itself(essentially partition args from lambdas)

By partition args and lambdas, we lose positional data, which are necessary to implement SqlUnparser and logical and physical formatting, leading to either the implementation itself having to implement them, adding some boilerplate, where today is automatically handled by ScalarUDF and SqlUnparser itself, or that the implementation save and expose positional data, that is them used by ScalarUDF and SqlUnparser to de-partition the args list, and them proceed normally.

Some closed system without support for udf with lambdas use a common positioning for lambdas and can just use them without having to store positional data nor lambdas and other args in the same ordered list:

Always last for spark, duckdb and snowflake:

transform(array(1), v -> v*2) -- spark
list_transform([1, 2, NULL, 3], lambda x: x + 1)  -- duckdb
transform([1, 2, 3], a INT -> a * 2) -- snowflake

Always first for clickhouse:

 arrayMap(x -> (x + 2), [1, 2, 3])

But I think that datafusion as an open and extendable system shouldn't stick to a single convention as different deployments may want to support one, another or both conventions.

Furthermore, my interest in lambda functions is to implement union manipulating udfs, which receives multiple lambdas interleaved with regular arguments. Consider an Union<'str' = Utf8, 'num' = Float64>, that could be transformed like this:

union_transform(
    union_value, 
    'str', str -> trim(str),
    'num', num -> num*2
)

where today the only way is with this(if/when support for union_value lands):

case union_tag(union_value)
    when 'str' then union_value('str', trim(union_extract(union_value, 'str')))
    when 'num' then union_value('num', union_extract(union_value, 'num') * 2)
end

For such cases, storing positional data or manually implementing logical and phyiscal formatting and SqlUnparser for every UDF is mandatory, as that positioning doesn't fit the simple conventions like always first or last.

This also omit positional info from the implementation. None lambda udf discussed so far, including union_transform, requires positional info for evaluation, only for logical and physical formatting and SqlUnparser, but I don't like the idea of blocking any future hypothetical udf that requires positional info for evaluation. A theoretical example is union_transform itself supporting a slight shorter syntax for lambdas that returns constants:

union_transform(
    union_value, 
    'str', str -> trim(str),
    'num', _num -> 0 -- lambda that returns constant/scalar
)

union_transform(
    union_value, 
    'str', str -> trim(str),
    'num', 0 -- shorter syntax, pass constant directly
)

That syntax does requires positional info for evaluation. It's really just an example and I don't plan to support it on my union_transform implementation

Instead, if we add a new LambdaFunction expression and LambdaUDF trait, the implementation could own both the lambdas and the other args, so that no partitioning would occur:

struct LambdaFunction { invocation: Box<dyn LambdaInvocation> } //invocation includes all args, both lambdas and non-lambdas
Expr::LambdaFunction(LambdaFuntion::new(ArrayTransform::new(all_args_including_lambdas)))

// instead of 
struct LambdaFunction { func: Arc<dyn LambdaUDF>, args: Vec<Expr> } // func owns only lambdas, and non lambdas are stored on args
Expr::LambdaFunction(LambdaFunction::new(ArrayTransform::new(lambda), args_without_lambdas))

Insights from Spark implementation

I'm don't have much experience with Spark and Scala, I only implemented a custom data source a few years ago, so if something below seems wrong, it probably is.

Spark includes both HigherOrderFunction trait and LambdaFunction class, similar to our current Expr::Lambda (and not the proposed LambdaFunction), as well as NamedLambdaVariable class, similar to ours Expr::LambdaVariable, all of them extending Expression, which in turn extends TreeNode

Some traversals branch directly on LambdaFunction expressions, which in my view support the idea of both keeping Expr::Lambda and exposing the lambda body on the TreeNode API:

Lambda Binder, (similar to ours BindLambdaVariable)
ResolveLambdaVariables
ColumnResolutionHelper
ColumnNodeToExpressionConverter
CheckAnalysis
ColumnNodeToProtoConverter

Some on branch LambdaVariable, which in my view support the ideia exposing of the lambda body on the TreeNode API:

SessionCatalog
HigherOrderFunction.functionForEval
NormalizePlan
TableOutputResolver

Some branch on HigherOrderFunction:

CheckAnalysis
Analyzer
ResolveLambdaVariables

Is true that some traversals branch on HigherOrderFunction expressions, therefore supporting the ideia of a new Expr::LambdaFunction, but since checking if a Expr::ScalarFunction invocation contains lambdas is simple as the code below, and can be made even simpler by adding a helper method contains_lambdas(&self) -> bool to ScalarFunction, I don't think it alone justifies adding a new Expr::LambdaFunction holding a ScalarFunction. If LambdaUDF trait is added then a new Expr::LambdaFunction variant has to be added anyway so this become a non-issue.

match expr {
    Expr::ScalarFunction(fun) if fun.args.iter().any(|arg| matches!(arg, Expr::Lambda(_))) => ...,
    // with helper
    Expr::ScalarFunction(fun) if fun.contains_lambdas() => ...
    ...
}

impl ScalarFunction {
    fn contains_lambdas(&self) -> bool {
        self.args.iter().any(|arg| matches!(arg, Expr::Lambda(_)))
    }

    // Other helpers could be added to like:
    fn lambda_args(&self) -> impl Iterator<Item = &LambdaFunction> {
        self.args.iter().filter_map(|arg| match arg {
            Expr::Lambda(l) => Some(l),
            _ => None
        })
    }
    
    fn non_lambda_args(&self) -> impl Iterator<Item = &Expr> {
        self.args.iter().filter(|arg| !matches(arg, Expr::Lambda(_)))
    }
}

Higher order functions, like ArrayTransform, receive both the regular arg as well as the lambda as regular Expressions, the lambda being expected to be a LambdaFunction:

case class ArrayTransform(
    argument: Expression,
    function: Expression)
  extends ArrayBasedSimpleHigherOrderFunction with CodegenFallback {

  override def dataType: ArrayType = ArrayType(function.dataType, function.nullable)

  override protected def bindInternal(
      f: (Expression, Seq[(DataType, Boolean)]) => LambdaFunction): ArrayTransform = {
    val ArrayType(elementType, containsNull) = argument.dataType
    function match {
      case LambdaFunction(_, arguments, _) if arguments.size == 2 =>
        copy(function = f(function, (elementType, containsNull) :: (IntegerType, false) :: Nil))
      case _ =>
        copy(function = f(function, (elementType, containsNull) :: Nil))
    }
  }

And it's instantiated like that:

        val param = NamedLambdaVariable("x", et, containsNull)
        val funcBody = insertMapSortRecursively(param)

        ArrayTransform(e, LambdaFunction(funcBody, Seq(param)))

So, regarding to resolving lambdas/partition lambdas from args, all args are owned by the implementation, and not by HigherOrderFunction, which is just a trait, therefore the lambas are indeed resolved, but so are the other args, and no partitioning occurs. I believe that our current approach would look like this in Spark:
trait Expression, implemented by class HigherOrderFunction which owns a trait HigherOrderFunctionImpl which is implemented by ArrayTransform, where the lambda and the arg would be stored in the class HigherOrderFunction, not in ArrayTransform. The Spark approach implemented in DF would be look the new LambdaFunction expression and LambdaUDF trait cited above, where implementation owns both the lambdas and the other args: Expr::LambdaFunction(LambdaFuntion::new(ArrayTransform::new(all_args_including_lambdas))), instead of Expr::LambdaFunction(LambdaFunction::new(ArrayTransform::new(lambdas), args_without_lambdas))

Finally, running this script on spark-shell returns the following output, showing both that lambda bodies are present in TreeNode traversals as well as LambdaFunction nodes:

val df = spark.sql("SELECT transform(array(1), v -> v*2)")
val transform = df.queryExecution.logical.expressions(0)

transform.foreach { node =>
  println(s"${node.nodeName}: ${node.toString}")
}

println(transform.treeString)

UnresolvedAlias: unresolvedalias(transform(array(1), lambdafunction((v * 2), v)))
UnresolvedFunction: transform(array(1), lambdafunction((v * 2), v))
UnresolvedFunction: array(1)
Literal: 1
LambdaFunction: lambdafunction((v * 2), v)
Multiply: (v * 2)
UnresolvedNamedLambdaVariable: v
Literal: 2
UnresolvedNamedLambdaVariable: v

unresolvedalias('transform('array(1), lambdafunction((lambda 'v * 2), lambda 'v, false)))
+- 'transform('array(1), lambdafunction((lambda 'v * 2), lambda 'v, false))
   :- 'array(1)
   :  +- 1
   +- lambdafunction((lambda 'v * 2), lambda 'v, false)
      :- (lambda 'v * 2)
      :  :- lambda 'v
      :  +- 2
      +- lambda 'v

In short, if we more or less agree with my points above, adding a new LambdaUDF just to not modify ScalarFunctionArgs, IMHO, doesn't seem worth the additional code that get's to be reviewed and maintained both in this PR (which is already big) and in subsequent PRs adding more lambdas UDFs like array_filter and array_fold

What do you think? I'm missing or misunderstanding something? Thanks!
And again, sorry for the delay

cc @keen85

Sorry for missing the update on this thread and thanks for the detailed analysis @gstvg!

The reasoning makes a lot of sense to me. And it's great to see your investigation on the Spark codebase.

I didn't think deeply enough about how the lambda function would need to interact with various parts in DataFusion, such as tree traversal, SQL unparser, and the general ability to preserve argument positions. So the LambdaUDF trait I mentioned indeed has many limitations.

Therefore, feel free to continue the ideas you are pursuing! And let's see if DataFusion maintainers who are familiar with this part of the codebase have anything to add for the approach.

Thanks @gstvg for driving this. I missed this PR and actually was advocating for lambda support in DataFusion roadmap. Let me explore this approach, in high level IMO it makes sense to introduce LambdaUDF trait so taking as example SELECT list_filter([1,2,3,4], x -> x % 2 = 0); it could look like

FunctionExpression
  name: list_filter
  children:
    1. ListExpression [1,2,3,4]
    2. LambdaExpression
         parameters: ["x"]
         body:
            ComparisonExpression (=)
              left:
                 ArithmeticExpression (%)
                    left: ColumnRef("x")
                    right: Constant(2)
              right:
                 Constant(0)

Then we likely need to Bind x to know what is it and finally rewrite it into something computable, perhaps using UDF of (Int) => Boolean and provide a bitmap as the result of valid indices and then extract valid indices from array. It might preserve SIMD execution.

It is great to see this effort is moving

Thanks @linhr! It most applies to having lambdas and args partitioned, omitting the body on TreeNode and removing Expr::Lambda. Changing the PR to just use a new Expr::LambdaFunction(Arc<dyn LambdaUDF>) where args are unpartitioned, body is exposed on TreeNode and Expr::Lambda is kept is okay, it wouldn't incur in the other cited points, but would make the PR bigger.

Thanks @comphead! The current representation looks somewhat similar to your suggestion, except that is a Expr::ScalarFunction/ScalarFunctionExpr instead of dedicated LambdaFunction, and that references to lambda parameters uses LambdaVariable instead of Column, containing an optional value that should be binded before execution:

FunctionExpression //regular ScalarFunction
  name: list_filter
  children:
    1. ListExpression [1,2,3,4]
    2. LambdaExpression // the added physical LambdaExpr / logical Expr::Lambda
         parameters: ["x"]
         body:
            ComparisonExpression (=)
              left:
                 ArithmeticExpression (%)
                    left: LambdaVariable("x", Field::new("", Int32, false), None) that after binding becomes
                          LambdaVariable("x", Field::new("", Int32, false), Some([1, 2, 3, 4]))
                    right: Constant(2)
              right:
                 Constant(0)

The reason why adding LambdaVariable is easier than trying to use Column is show on the outdated sections on this PR description and on the first comment.

If the added code for a LambdaUDF is considered worthwhile to not modify ScalarUDF and/or because it's a better representation, while keeping others things unchanged, I'm okay with it. But since you mentioned vectorized execution, I just want to make sure that we are not moving to LambdaUDF only for performance because, contrary to Spark, where the LambdaVariable value is set and re-set for every row, which would obviously be terrible for a vectorized engine, here it is only set once per batch, and so can perform as fast as regular expressions. The LambdaVariable value being a ColumnarValue is just a niche optimization for scalar evaluation, we can change it to support ArrayRef only. Native performance has been a goal since the beginning, and a list_filter implementation based on this can be easily vectorized. As an example, #17220 got vectorized execution while also based on ScalarUDF and regular physical expressions, the lambda body is evaluated once per batch and it's output used to filter the list values with the filter kernel, and then the offsets are adjusted:

    let filter_array = lambda.evaluate(&batch)?;
    let ColumnarValue::Array(filter_array) = filter_array else {
        return exec_err!("array_filter requires a lambda that returns an array of booleans");
    };

    let filter_array = as_boolean_array(&filter_array)?;
    let filtered = filter(&values, filter_array)?;

    for row_index in 0..list_array.len() {
        if list_array.is_null(row_index) {
            // Handle null arrays by keeping the offset unchanged
            offsets.push_length(0);
            continue;
        }
        let start = value_offsets[row_index];
        let end = value_offsets[row_index + 1];
        let num_true = filter_array
            .slice(start.as_usize(), (end - start).as_usize())
            .true_count();
        offsets.push_length(num_true);
    }
    let offsets = offsets.finish();
    let list_array = GenericListArray::<OffsetSize>::try_new(
        Arc::clone(field),
        offsets,
        filtered,
        nulls.cloned(),
    )?;

    Ok(Arc::new(list_array))

One thing that can be optimized is use Buffer::count_set_bits_offset instead of BooleanArray::slice and then true_count to avoid 1-2 Arc::clone per loop iteration. There's also a faster way to adjust the offsets when the average sub-list length is small, something like 64 or less, but it's bigger and a bit more intricate, so I choose implement array_transform instead to make the PR smaller. A reasonable fast array_fold is even more complicated, but still implementable in any lambda approach discussed here.

array_transform here is also vectorized: the transforming lambda body is evaluated only once per batch/invoke/evaluate call.

Then we likely need to Bind x to know what is it and finally rewrite it into something computable, perhaps using UDF of (Int) => Boolean

I believe that by using LambdaVariable, planning the lambda body into a physical expr as usual is enough to compute it, without having to use UDFs. One thing that we must take into account is column capture, which is already fully supported here, including shadowing and projection pushdown(columns from the input plan that only appear within a lambda body). As an example, #17220 does not support that, and if tried to support it while using Expr::Column for lambda variables and passing those variables in the RecordBatch to PhysicalExpr::evaluate, would likely get the same problems I got in my first approach, which are also show in the outdated sections of the description and of the first comment here. That's why LambdaVariable exists. Supporting lambdas without capture support is relatively easy, regardless of the approach.

@rluvaton I see you also were not fond of adding physical expr's to ScalarFunctionArgs on #17220. Any comments here?

Finally, how would the new LambdaUDF trait works?
I believe it should look like ScalarFunction?

struct LambdaFunction {
    args: Vec<Expr>,
    fun: Arc<dyn LambdaUDF>,
}

enum Expr {
    ...
    Expr::LambdaFunction(LambdaFunction),
}

It also can be made more like Spark:

trait LambdaInvoker {
    ...
    fn invoke(&self, args: Vec<Expr>) -> Box<dyn LambdaInvocation>;
}

enum Expr {
    ...
    Expr::LambdaFunction(Box<dyn LambdaInvocation>),
}

I believe LambdaUDF/LambdaInvocation trait would look like ScalarUDFImpl except for the lambdas_parameters method, ReturnFieldArgs in return_field_from_args and LambdaFunctionArgs in the invoke method, or any of you imagine more fundamental changes?

Thanks!

thanks a lot for this PR.

Here are my thoughts:

Future features that the API should support without or minimal breaking changes

Couple of things to make sure we support or have a way to add them in the future without breaking changes:

1. Support Map, (Large)List, (Large)ListView, FixedSizeList as the input for lambda
2. multiple lambdas in a single expression, for example map_key_value(some_map_col, map_key_lambda, map_value_lambda) and each lambda gets a different variables
3. Lambda expression that access columns that are not in the list itself
   so I can do the following:

   | year |    grades      |
   |------|----------------|
   | 1998 | [1, 2, 3]      |
   | 1999 | [4, 99, 5, 10] |
   | 2000 | [6, 0, null]   |
   

   array_transform(grades, x -> if year <= 1990 then x * 10 else x)
4. optional arguments for lambda, for example the index of the item in the list
   the optional here is important as I want to avoid creating that input if I don't need to.
5. Nested lambda expressions: array_transform(matrix, x -> array_transform(x, y -> y * 2))

Things that every lambda implementation would need to handle

1. replace null lists that is not empty underneath, with nulls with empty list underneath.
   consider the following when the expression can fail: array_transform(list, x -> 1 / x)
   with this input

   fn get_list() -> GenericListArray<i32> {
     GenericListArray::new(
       Arc::new(Field::new_list_field(DataType::Int8, false)),
       OffsetBuffer::<i32>::from_lengths(vec![2, 2, 1]),
       Arc::new(Int8Array::from(vec![1, 2, 0, 3, 4])),
       Some(NullBuffer::from(&[true, false, true])),
     )
   }

   which the second list is null but the underlying value is [0, 3] which if we run the transform on 0 on it, it will fail with division by zero.

   I have a lot of helpers to cleanup the nulls BTW

2. sliced list and not computing values outside the sliced data

3. fixed size list (which is different than list as you can't change the underlying nulls to be empty list) with a null list and the underlying values can cause failures, consider the expression array_transform(list, x -> 1 / x) on this input:

fn get_list() -> FixedSizeListArray<i32> {
    FixedSizeListArray::new(
       Arc::new(Field::new_list_field(DataType::Int8, false)),
       3,
       Arc::new(Int8Array::from(vec![
           1, 2, 3,
           0, 1, 1,
           4, 5, 6
       ])),
       Some(NullBuffer::from(&[true, false, true])),
    )
}

In every case here, we should answer:

1. how would the user handle that?
2. how could we make it easier for the user?
3. How would we help them avoid forgetting handling?

I'm conflicted on whether we should fix case 1 for them as it is costly and the user might have prior knowledge to avoid that.

Why a new trait LambdaUDFImpl instead of adding functions on ScalarUDF

I think having a new LambdaUDFImpl is better than adding functions on existing ScalarUDF because:

1. The ScalarUDF trait will not grow too much and make implementing regular scalar UDFs easier or lambda overwhelming
2. what if we need to add a required function but only for lambda, we can add it on the new trait with ease and we won't need to do some weird stuff
   to avoid breaking changes.
3. Less ambiguity on the API.

Implementation specific:

I want to keep the simplicity of ScalarUDF which means that in order to evaluate a lambda expression I don't need to construct stuff, only need to provide the input and maybe some options for future use.

Thanks @rluvaton and @gstvg , its nice you mentioned array_transform, the tricky part for this function is its return type depends on lambda

array_transform(array<T>, function<T, U>) -> array<U>

I want to keep the simplicity of ScalarUDF which means that in order to evaluate a lambda expression I don't need to construct stuff, only need to provide the input and maybe some options for future use.

Right, on high level it could be like

pub struct LambdaExpr {
    /// Parameter names/types already resolved
    pub param_types: Vec<DataType>,

    /// Expression body, what needs to be evaluated, this thing potentially can be UDF
    pub body: Arc<dyn PhysicalExpr>,
}

Impl

impl LambdaExpr {
    pub fn new(
        param_types: Vec<DataType>,
        body: Arc<dyn PhysicalExpr>,
    ) -> Self {
        Self { param_types, body }
    }

    /// Evaluate lambda over provided arrays
    pub fn evaluate_with_args(
        &self,
        args: Vec<ArrayRef>,
    ) -> Result<ArrayRef> {
        // Build synthetic schema
        let fields: Vec<Field> = self.param_types
            .iter()
            .enumerate()
            .map(|(i, dt)| Field::new(format!("arg{}", i), dt.clone(), true))
            .collect();

        let schema = Arc::new(Schema::new(fields));

        let batch = RecordBatch::try_new(schema, args)?;

        self.body.evaluate(&batch)   // this where our UDF would be called
    }
}

So for example x -> x + 1 we need to parse expression and create our Lambda, so we need to modify parser to get structures below from user defined code and there is an existing ticket apache/datafusion-sqlparser-rs#1273

// Parameter x at column 0
let x = Arc::new(ColumnExpr::new(0));

// Literal 1
let one = Arc::new(LiteralExpr::new(
    ScalarValue::Int32(Some(1))
));

// x + 1
let body = Arc::new(BinaryExpr::new(
    x,
    one,
    Operator::Add,
));

// Lambda(x) -> x + 1
let lambda = LambdaExpr::new(
    vec![DataType::Int32],
    body,
);

and call it from caller built in function

fn array_transform(
    list_array: &ListArray,
    lambda: &LambdaExpr,
) -> Result<ListArray> {

    let values = list_array.values().clone();

    // evaluate lambda on flattened child array
    let transformed =
        lambda.evaluate_with_args(vec![values])?;

    Ok(ListArray::new(
        list_array.data_type().clone(),
        list_array.offsets().clone(),
        transformed,
        list_array.nulls().cloned(),
    ))
}

@rluvaton @comphead

I updated the PR to a LambdaUDF trait based implementation. It added 1300 lines, totaling 3000, mostly boilerplate from ScalarUDF including a lot of documentation. Was that on the range you were expecting?
Based on the ScalarUDF docs stating that it exists to maintain backwards compatibility with an older API, I included only a LambdaUDF trait and not a struct LambdaUDF + trait LambdaUDFImpl pair to not make even bigger. There's a few small things missing that I want to implement tomorrow, and if all of you are okay with the results, I want to open this to review. My only concern is that the PR size may delay the review, as in #17220 (comment), but since IMHO this PR is simpler, I hope it will be that long to review, despite the size

1. The ScalarUDF trait will not grow too much and make implementing regular scalar UDFs easier or lambda overwhelming
2. what if we need to add a required function but only for lambda, we can add it on the new trait with ease and we won't need to do some weird stuff to avoid breaking changes.
3. Less ambiguity on the API.

Yeah, I think that 2 is the main point. The previous ScalarUDF based approach only added a single method to the trait that already contains 20, it didn't require any change for non-lambda udfs (this would be unacceptable), and compared to the actual LambdaUDF based, lambda UDFs only required 2 additional lines of code per implementation, but, while I can't think of any ... all my previous counter-arguments can fall apart in the future with a single requirement change. There's upfront cost in review time and time to merge, but also more room to work in the future

Support Map, (Large)List, (Large)ListView, FixedSizeList as the input for lambda

Currently any type is accepted and passed to to the implementation to derive the parameters from it, usually the inner values of a list, but I want to work with unions too, for example. But as of now array_transform doesn't handle ListView's (I thought there's no plans to support it overall?). Since the ListView values may contain unreferenced values, should it be compacted, or casted to a regular compacted List? And if so, return the transformed List or cast it to ListView?

multiple lambdas in a single expression, for example map_key_value(some_map_col, map_key_lambda, map_value_lambda) and each lambda gets a different variables

Supported: LambdaFunctionArgs.args can hold multiple lambdas

    fn invoke_with_args(&self, args: LambdaFunctionArgs) -> Result<ColumnarValue> {
        let [list_value, lambda] = take_function_args(self.name(), &args.args)?;

        let (ValueOrLambda::Value(list_value), ValueOrLambda::Lambda(lambda)) =
            (list_value, lambda)
        else {
            return exec_err!(...

Since the beginning I worked to support multiple lambdas to implement union manipulating functions like this:

union_transform(
    union_value, 
    'str', str -> trim(str),
    'num', num -> num*2,
    'bool', bool -> NOT bool,
)

Lambda expression that access columns that are not in the list itself

Already supported and tested (I call it column capture through the PR and comments. Please ignore the comment above the test, I'm going to remove it)

CREATE TABLE t as SELECT 1 as n;
query ?
SELECT array_transform([1, 2], (e) -> n) from t;
----
[1, 1]

optional arguments for lambda, for example the index of the item in the list
the optional here is important as I want to avoid creating that input if I don't need to.

Also supported

        // use closures so that bind_lambda_variables evaluates only the params that are actually referenced
        // avoiding unnecessary computations
        let values_param = || Ok(Arc::clone(list_values));
        let indices_param = || elements_indices(&list_array);

        let binded_body = bind_lambda_variables(
            Arc::clone(&lambda.body),
            &lambda.params,
            &[&values_param, &indices_param],
        )?;

*after a refactor, bind_lambda_params actually eager evaluated all params, but I'll fix that

Nested lambda expressions

Supported

SELECT array_transform(t.v, (v1, i) -> array_transform(v1, (v2, j) -> array_transform(v2, v3 -> j)) ) from t;
----
[[[0, 1], [0]], [[0]], [[]]]

array_transform, the tricky part for this function is its return type depends on lambda

LambdaUDF contains a method lambdas_arguments where the implementation must return the type of the parameters of all it's lambdas when evaluated with a given set of values. Then DF uses this info to compute the type of the lambdas and pass it on return_field_from_args, so the implementation can easily compute it's return type. For example the expr array_transform([1, 2], v -> repeat(v, 'a')), lambdas_arguments would receive [ValueOrLambda::Value(List(Int32)), ValueOrLambda::Lambda] and should return vec![None, Some(vec![DataType::Int32])]. Then return_field_from_args would be called with [ValueOrLambda::Value(List(Int32)), ValueOrLambda::Lambda(Utf8)], where the implementation would need just to return List(Utf8).

pub struct LambdaReturnFieldArgs<'a> {
    /// The data types of the arguments to the function
    ///
    /// If argument `i` to the function is a lambda, it will be the field returned by the
    /// lambda when executed with the arguments returned from `LambdaUDF::lambdas_parameters`
    ///
    /// For example, with `array_transform([1], v -> v == 5)`
    /// this field will be `[Field::new("", DataType::List(DataType::Int32), false), Field::new("", DataType::Boolean, false)]`
    pub arg_fields: &'a [ValueOrLambdaField],
    /// Is argument `i` to the function a scalar (constant)?
    ///
    /// If the argument `i` is not a scalar, it will be None
    ///
    /// For example, if a function is called like `my_function(column_a, 5)`
    /// this field will be `[None, Some(ScalarValue::Int32(Some(5)))]`
    pub scalar_arguments: &'a [Option<&'a ScalarValue>],
}

/// A tagged Field indicating whether it correspond to a value or a lambda argument
#[derive(Clone, Debug)]
pub enum ValueOrLambdaField {
    /// The Field of a ColumnarValue argument
    Value(FieldRef),
    /// The Field of the return of the lambda body when evaluated with the parameters from LambdaUDF::lambda_parameters
    Lambda(FieldRef),
}

    fn return_field_from_args(
        &self,
        args: datafusion_expr::LambdaReturnFieldArgs,
    ) -> Result<Arc<Field>> {
        let [ValueOrLambdaField::Value(list), ValueOrLambdaField::Lambda(lambda)] =
            take_function_args(self.name(), args.arg_fields)?
        else {
            return exec_err!(
                "{} expects a value follewed by a lambda, got {:?}",
                self.name(),
                args
            );
        };

        //TODO: should metadata be copied into the transformed array?

        // lambda is the resulting field of executing the lambda body
        // with the parameters returned in lambdas_parameters
        let field = Arc::new(Field::new(
            Field::LIST_FIELD_DEFAULT_NAME,
            lambda.data_type().clone(),
            lambda.is_nullable(),
        ));

        let return_type = match list.data_type() {
            DataType::List(_) => DataType::List(field),
            DataType::LargeList(_) => DataType::LargeList(field),
            DataType::FixedSizeList(_, size) => DataType::FixedSizeList(field, *size),
            _ => unreachable!(),
        };

        Ok(Arc::new(Field::new("", return_type, list.is_nullable())))
    }