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

Benchmark	Spark (ms)	Comet Native (ms)	Comet Codegen (ms)	Native vs Spark	Native vs Codegen
int literal	5156	1408	5234	3.7x	3.7x
capture outer column	6582	1449	5147	4.5x	3.6x
compound predicate (AND / range)	8659	1498	7413	5.8x	4.9x
arithmetic expression in lambda	8826	1545	7494	5.7x	4.8x
string length predicate	20524	2579	17323	8.0x	6.7x
string equality comparison	12657	3051	8896	4.1x	2.9x
array with nulls (IS NOT NULL check)	8703	1798	6855	4.8x	3.8x
nested array (size check)	3752	1028	2121	3.6x	2.1x
chained filters (pipeline)	10115	1898	15495	5.3x	8.2x
short arrays	933	203	627	4.6x	3.1x
large arrays	63915	14675	54114	4.4x	3.7x

[andygrove]

Note on this review: this was generated by an LLM (Claude Code) at my request while I worked through a review backlog. I have not verified the individual findings myself. Please treat everything below as suggestions to evaluate rather than as authoritative review feedback, and push back on anything that is wrong or already handled.

Moving lambda evaluation into DataFusion instead of paying a JNI call per batch is clearly the right direction, and the design here is thoughtful: resolving NamedLambdaVariable by Spark exprId rather than by name sidesteps shadowing, and the wrapper that keeps unused lambda parameters visible in children() so projection compaction stays consistent with the runtime layout is a subtle detail well handled.

Four things.

The configs are read once per JVM, not per session

case class CometHighOrderFunction[T <: HigherOrderFunction](name: String)
    extends CometExpressionSerde[T] {
  private val nativeHofEnabled = CometConf.COMET_EXEC_HIGHER_ORDER_FUNCTION_NATIVE_ENABLED.get()
  private val codegenEnabled = CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.get()

CometArrayFilter is an object extending this case class, so both vals are evaluated once when the object is first initialized, and never again. That means spark.comet.exec.higherOrderFunction.native.enabled is effectively frozen at whatever it was on first use, and withSQLConf(...) in a test or a per-session override in production will not take effect.

Other serdes read configs inside getSupportLevel / convert precisely for this reason. Could these become defs, or be read from the conf passed into the conversion? This is also going to make the feature untestable with withSQLConf, which is worth checking against the existing tests: if a test appears to exercise both paths, it may only be exercising whichever ran first.

Defaulting the native path to true

A new native lambda execution path with a scope stack, exprId resolution, and projection compaction, enabled by default in its first release. Given how much new machinery this is, would false be safer for one release? The fallback chain (native, then dispatcher, then Spark) is already in place, so users who want it can opt in and the risk of a subtle lambda-scoping bug reaching everyone is much lower.

The module doc no longer matches with_scope

lambda.rs's header says the planner needs "(2) A drop-guard that pops a scope on any exit path (?, panic-safe)", but with_scope's own comment says "The pop happens on both the Ok and Err paths, this replaces the earlier RAII guard". with_scope is not panic-safe: a panic inside f leaves the scope on the stack.

Since PhysicalPlanner catches panics across the FFI boundary and could in principle keep planning, that could leave a stale scope that resolves a variable it should not. Either restore the guard or fix the header so it does not promise panic safety it no longer provides.

No performance number

The stated motivation is removing a per-batch JNI call. What does that buy for array_filter over a realistic batch? A microbenchmark or an end-to-end comparison against the dispatcher path would justify the complexity, and would also tell reviewers whether extending this to transform, exists, and aggregate is worth doing.

[erikbogado-nstech]

Posting benchmark numbers for this PR from the suite I mentioned on the
dev list (https://github.com/ErikBPF/ndc — derived from TPC-H over
nested layouts; answers pinned, row parity enforced per query).

Setup: Spark 4.1.3 on a single 28-core host with local NVMe. 24
queries × 3 runs, page cache dropped between runs, medians. Built
main (75fdddc) and this PR's head from source with identical flags
(-Prelease jars, packaged native lib). Caveat: the PR head no longer
rebases cleanly over native/core/src/execution/planner.rs, so numbers
are directional w.r.t. the diff.

Results (Parquet, speedup vs vanilla Spark, 24 queries):

	main	this PR
built-in (scan+agg)	~2.3x	~2.2x
depth 1–8	~1.1x	~1.1x
extended	~1.08x	~1.31x
total	1.26x	1.42x

The two queries that lose to vanilla on main recover largely here:

• e21_nested (correlated quantifiers as boolean HOF filter, no
  explode): 2735 → 567 ms (−79%)
• e19_nested_array (IN-lists + arithmetic inside aggregate): −33%

A few non-lambda extended queries regress on the PR head (e14 +54%,
e07 +41%, e13 +26%) — base drift or lambda-planner adoption, worth a
look on rebase.

Raw records (per-query medians, native/fallback share, parity flags)
happy to share.

[kazantsev-maksim]

Removed LambdaParamsCapture

The LambdaParamsCapture / pin_unused_params wrapper in the native code is gone entirely. It existed to anchor unused lambda parameters in the expression tree so DataFusion's optimizer could not prune them and break the runtime batch layout.

This workaround became unnecessary after apache/datafusion#24162 landed in DataFusion 55: LambdaExpr now computes used_param_indices() itself, and LambdaArgument::new pushes only the parameters actually referenced by the body into the evaluation batch (captures ++ used_params, in declaration order). The runtime layout contract is now enforced on the DataFusion side, so the planner assigns lambda variables their declared positions and DataFusion compacts them - no anchoring wrapper needed. Nested-lambda scoping (the exprId-keyed scope stack in lambda.rs) is unaffected and still handles shadowing.

Verified: three levels of nesting, an inner HOF whose value argument is an outer lambda variable, sibling nested HOFs, an inner lambda referencing an outer variable, and multi-param lambdas with unused parameters all match Spark.

Benchmark results

Simple benchmark result (Apple M1 Pro, OpenJDK 17.0.19, 2 iterations, single run on final code; dispatch-path selection verified via serde logging):

Benchmark	Spark (ms)	Comet Native (ms)	Comet Codegen (ms)	Native vs Spark	Native vs Codegen
int literal	5156	1408	5234	3.7x	3.7x
capture outer column	6582	1449	5147	4.5x	3.6x
compound predicate (AND / range)	8659	1498	7413	5.8x	4.9x
arithmetic expression in lambda	8826	1545	7494	5.7x	4.8x
string length predicate	20524	2579	17323	8.0x	6.7x
string equality comparison	12657	3051	8896	4.1x	2.9x
array with nulls (IS NOT NULL check)	8703	1798	6855	4.8x	3.8x
nested array (size check)	3752	1028	2121	3.6x	2.1x
chained filters (pipeline)	10115	1898	15495	5.3x	8.2x
short arrays	933	203	627	4.6x	3.1x
large arrays	63915	14675	54114	4.4x	3.7x

The native path is 3.6-8.0x faster than vanilla Spark and 2.1-8.2x faster than the JVM codegen dispatch path across all scenarios. The codegen dispatch path (running Spark's own lambda evaluation inside the Comet kernel) is on par with or slower than vanilla Spark on lambda-heavy queries (chained filters: 0.7x), confirming the per-batch JNI and row-wise evaluation overhead this PR removes.

@comphead @andygrove Could you please take another look?