feat: Lambda function support from DataFusion, illustrated with array_filter#4744
feat: Lambda function support from DataFusion, illustrated with array_filter#4744kazantsev-maksim wants to merge 90 commits into
Conversation
This reverts commit 768b3e9.
# Conflicts: # native/proto/src/proto/expr.proto
|
@andygrove @comphead, thanks for the review. I tried to address your comments. Could you please take another look when you have a chance? |
| query | ||
| SELECT filter(filter(arr, x -> x < threshold), y -> y > 0) FROM test_array_filter | ||
|
|
||
| query | ||
| SELECT filter(filter(arr, x -> x < threshold), y -> y > 0) FROM test_array_filter |
There was a problem hiding this comment.
These two tests seem identical?
| self.create_high_order_function_expr(hof, input_schema) | ||
| } | ||
| ExprStruct::NamedLambdaVariable(nlv) => { | ||
| let idx = input_schema.index_of(&nlv.name).map_err(|_| { |
There was a problem hiding this comment.
Lambda variables here resolve by name via input_schema.index_of, which returns the first field of that name. That's only safe when the name is unique across the combined schema. Two cases concern me: a lambda variable whose name collides with an input column (e.g. SELECT x, filter(arr, x -> x > 2) FROM t), and nested lambdas that reuse a name (filter(arr, x -> exists(arr2, x -> ...))). In both, index_of returns the wrong field and we'd silently read the wrong column, whereas Spark disambiguates by ExprId. Could we add tests for both? If they diverge, resolving by position or uniquifying the names in the serde would close the gap.
There was a problem hiding this comment.
Correction to my comment above: I checked this out locally and the single-level column-collision example I gave (filter(arr, x -> x > 2) with a column also named x) actually works fine. DataFusion's LambdaArgument/merge_captures_with_variables machinery handles the flat case correctly, so please disregard that example.
The real problem is broader though: any higher-order function nested inside another lambda's body resolves lambda variables to the wrong index. In create_lambda_expr the body schema is input_schema.fields() ++ this lambda's args, but for an inner lambda input_schema already contains the outer lambda's appended arg, so the index no longer lines up with what LambdaExpr's projection remapping expects. It shows up two ways.
Distinct variable names produce a hard native crash:
SELECT filter(outer_arr, x -> size(filter(inner_arr, y -> y > 15)) > 0) FROM t
-- CometNativeException: Field of physical LambdaVariable with index 0 doesn't match
-- batch field during evaluation Field { "y": Int32 } != Field { "x": Int32 }Reusing the name silently returns wrong results:
SELECT filter(outer_arr, x -> size(filter(inner_arr, x -> x > 15)) > 0) FROM t
-- Spark: [[1, 2, 3]] Comet: [[]]With t as VALUES (array(1,2,3), array(10,20,30)). Both go native (CometProject in the plan). The existing filter(filter(arr, x -> ...), y -> ...) test passes because the inner filter is a value argument, not nested inside a lambda body, so it never actually nests.
This might be worth an explicit unsupported-shape guard (fall back to codegen/Spark when a lambda body itself contains a higher-order function) until nested lambda scoping is handled, so it doesn't crash or silently miscompute. Happy to share the SQL-file repros.
There was a problem hiding this comment.
Really appreciate you surfacing those edge cases! I'd be super grateful for any additional test cases you've got.
There was a problem hiding this comment.
Added more test cases
| -- Config: spark.comet.exec.scalaUDF.codegen.enabled=false | ||
|
|
||
| query expect_fallback(The array_filter function in DataFusion is limited to one lambda parameter) | ||
| SELECT filter(array(1, 2), (x, i) -> i > 0) |
There was a problem hiding this comment.
we may want to extend it in DF
| SELECT filter(filter(arr, x -> x < threshold), y -> y > 0) FROM test_array_filter | ||
|
|
||
| query | ||
| SELECT filter(arr, (x, i) -> i > 0) FROM test_array_filter |
There was a problem hiding this comment.
this tests passes but similar one should fallback in fallback file? 🤔
|
I reverted the label to Draft; I need a bit of time to figure out the handling of nested lambda functions. |
# Conflicts: # native/core/src/execution/planner.rs
|
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 |
# Conflicts: # spark/src/main/scala/org/apache/comet/serde/arrays.scala # spark/src/test/resources/sql-tests/expressions/array/array_filter.sql
3 participants
Which issue does this PR close?
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?
Simple benchmark result: