feat: Lambda function support from DataFusion, illustrated with array_filter

[kazantsev-maksim]

Which issue does this PR close?

• N/A

Rationale for this change

Running higher-order functions through JVM codegen is expensive: each batch incurs a JNI call into Spark's own implementation. Moving the lambda evaluation into the native DataFusion engine removes that overhead and brings the plan closer to fully native execution.

What changes are included in this PR?

• Protobuf (native/proto/src/proto/expr.proto) - Added three new messages: HigherOrderFunc (function name + value arguments + lambdas), LambdaFunction (body + arguments), and NamedLambdaVariable (name, type, nullable, expr_id). Added high_order_func (71) and named_lambda_variable (72) fields to Expr.

• Lambda Infrastructure & Scope Management - new lambda module: Introduced native/core/src/execution/lambda.rs to manage nested lambda variable scopes. This ensures that NamedLambdaVariables are correctly resolved by their Spark exprId, preventing name shadowing or column collisions. Optimizer Anchoring: Implemented LambdaParamsCapture (with a helper factory pin_unused_params). This is a critical mechanism to prevent DataFusion's optimizer from pruning "unused" lambda parameters. Since the runtime expects a specific batch structure, this wrapper "anchors" the parameters in the expression tree to maintain index consistency with the physical plan.

• Physical Planner Enhancements - HOF Planning: Extended PhysicalPlanner to support HigherOrderFunc expressions. It now includes logic to: Plan input value expressions. Query the UDF contract to resolve lambda parameter field types. Recursive plan the lambda body under the scope of created parameters.
  Variable Resolution: Added support for mapping NamedLambdaVariable protobuf definitions to physical LambdaVariable expressions, correctly binding them to the resolved indices in the lambda parameter schema.

• Infrastructure & Helpers - UDF Registration Helper: Added create_comet_hof_func in a new module comet_high_order_funcs.rs to simplify fetching supported HOFs from the DataFusion FunctionRegistry.
  Module Exposure: Updated native/core/src/execution/mod.rs to expose the new lambda-related modules to the rest of the core crate.

• Spark — serialization (CometHighOrderFunction.scala, QueryPlanSerde.scala, arrays.scala). New generic serializer CometHighOrderFunction[T] that converts a Spark HigherOrderFunction (along with LambdaFunction and NamedLambdaVariable) into protobuf. CometArrayFilter now extends this serializer: when spark.comet.exec.scalaUDF.codegen.enabled is disabled it takes the new native path, otherwise the old behavior is preserved (including the fast-path for array_compact).

How are these changes tested?

• Added new sql tests
• Added new benchmark test

Simple benchmark result:

[kazantsev-maksim]

I reverted the label to Draft; I need a bit of time to figure out the handling of nested lambda functions.

[kazantsev-maksim]

LambdaParamsCapture: This wrapper acts as a structural anchor. By forcing "unused" parameters to be reported as mandatory children() of the expression node, it effectively hides them from the pruning logic. This tricks the optimizer into treating them as "in-use," preventing it from removing or re-indexing them, and ensures that the physical index layout remains perfectly synchronized with the memory buffer provided by the runtime.

CC: @comphead

[comphead]

Thanks @kazantsev-maksim CI is green, I'll give it another look this week

[kazantsev-maksim]

We will be able to simplify the logic after merging this PR: apache/datafusion#23660

CC @comphead

[andygrove]

Thanks for the continued work on this. The nested-lambda scoping is genuinely fixed. I rebuilt the branch locally and re-ran every repro from my earlier round, plus harder ones: three levels of nesting, an inner HOF whose value argument is the outer lambda variable, two sibling nested HOFs in one body, and an inner lambda referencing the outer variable. All of them match Spark now. The exprId-keyed scope stack plus pin_unused_params is doing real work.

The three-tier config also reads much better than the earlier reuse of scalaUDF.codegen.enabled.

I did find one new issue that I think needs to be resolved before merge, plus a couple of smaller things. Details inline.

Disclosure: I used an LLM to help review this PR and to draft this comment. I built the branch and ran the queries below myself, and the results reported are from those runs.

One more, which I can't attach inline because the file isn't in the diff. docs/source/user-guide/latest/expressions.md line 366 still reads:

filter | ✅ | General lambda routed through the JVM codegen dispatcher; the array_compact form runs natively

That is no longer accurate for the common case now that the native path is the default. This page is hand-maintained rather than generated by GenerateDocs, so it needs updating here. The per-version pages under compatibility/expressions/ are generated at release time, so there is nothing to hand-edit there.

[andygrove]

I hit a case where the native path is taken but the lambda variable never reaches the evaluated body, so the predicate is false for every element. This is on fully default configuration, and it returns wrong results rather than erroring.

CREATE TABLE t(sarr array<string>, a array<int>, b array<int>) USING parquet;
INSERT INTO t VALUES (array('abc','xyz','a1'), array(1,2,3), array(10,20,30));

SELECT filter(sarr, x -> x rlike '^a') FROM t;
-- Spark: [abc, a1]   Comet: []

SELECT filter(a, x -> exists(b, y -> y > x)) FROM t;
-- Spark: [1, 2, 3]   Comet: []

SELECT filter(a, x -> array_max(transform(b, y -> y + x)) > 31) FROM t;
-- Spark: [2, 3]      Comet: []

All three plan to CometProject and return empty arrays. Setting spark.comet.exec.higherOrderFunction.native.enabled=false makes all three correct, which points at the new native path.

My reading of the mechanism: highOrderFunction2Proto calls exprToProtoInternal on the lambda body, and when the body contains rlike, a non-native nested higher-order function like exists or transform, or anything else that lands in CometScalaUDF.emitJvmCodegenDispatch, that call succeeds and emits a JvmScalarUdf proto. So hofProto.isDefined is true and we commit to native. But emitJvmCodegenDispatch binds against AttributeReferences only, and a NamedLambdaVariable is not one, so the per-element value never reaches the compiled kernel.

Could nativeUnsupportedReason walk the lambda body and decline the native path when it finds a subexpression that will route through the codegen dispatcher, including nested HigherOrderFunctions? Then convert degrades to codegen the way the config doc describes.

Worth SQL tests for these shapes as well. One thing to watch: filter(a, x -> exists(b, y -> y > 15)) passes, because nothing crosses the lambda boundary there. The capture has to be meaningful for the test to catch this.

[andygrove]

This replaces the parent's reasons rather than adding to them, so lambda functions must be LambdaFunction and lambda arguments must be NamedLambdaVariables from CometHighOrderFunction.getUnsupportedReasons() disappear from the generated compatibility docs for filter. Could this be super.getUnsupportedReasons() ++ Seq(UNARY_FUNCTION_EXPECTED)?

[kazantsev-maksim]

It's my mistake, fixed, thanks.

[andygrove]

The benchmark covers one shape only: x -> x > 2 over int arrays of size 100. Since the stated rationale is avoiding the per-batch JNI call, it would be more informative to also cover a predicate that captures an outer column and a string-element case, so we can see whether the win holds across the shapes people actually write.

[kazantsev-maksim]

I have updated CometArrayFilterBenchmark to cover a comprehensive set of shapes, including:

• Outer column capture (x -> x > threshold)
• String predicates (length(x) > 10, string equality)
• Compound boolean logic & arithmetic inside lambda
• Nullable element propagation (IS NOT NULL)
• Nested arrays (array<array<int>>)
• Chained filters pipeline & large SIMD-friendly arrays

Benchmark Results (Apple M1 Pro, 4M rows)

Benchmark Case	Spark (ms)	Comet Native (ms)	Comet Codegen (ms)	Relative (vs Spark)
int literal	4961	1083	1129	4.6x
capture outer col	5706	1110	1093	5.1x
compound predicate	7939	1123	1125	7.1x
arithmetic in lambda	8168	1227	1227	6.7x
string length	20797	2172	2144	9.6x
string equality	11769	2641	2639	4.5x
array with nulls	7897	1486	1488	5.3x
nested array	3705	973	971	3.8x
chained filters	9388	1587	1613	5.9x
short arrays	942	203	198	4.6x
large arrays (size 1000)	58994	11071	11317	5.3x

Comet consistently delivers 3.8x–9.6x speedups across all shapes. The native DataFusion path matches and occasionally edges out JVM codegen performance.