feat: Lambda function support from DataFusion, illustrated with array_filter - #4744
kazantsev-maksim wants to merge 129 commits into
Conversation
This reverts commit 768b3e9.
# Conflicts: # native/core/src/execution/mod.rs # native/core/src/execution/planner.rs
andygrove
added
enhancement
Removed LambdaParamsCaptureThe This workaround became unnecessary after apache/datafusion#24162 landed in DataFusion 55: 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 resultsSimple benchmark result (Apple M1 Pro, OpenJDK 17.0.19, 2 iterations, single run on final code; dispatch-path selection verified via serde logging):
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? |
# Conflicts: # native/core/src/execution/planner.rs
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Reviewed 95a43027916c3c7eb8206640c5997420091f57dd against 58ab5f618e1e715dee06165424672fdd820cafe4. I found two P2 correctness regressions: guarded lambda expressions can throw during serialization, and the native filter evaluates scalar predicates when the arrays contain no elements. Details and reproductions are inline. Both should be addressed before merging.
The revised dynamic configuration and JVM-dispatch guard passed the focused checks, including the earlier regex/exists/transform cases.
Validation:
- The root Maven reactor
test-compilepassed with JDK 17 and Spark 4.1.3. - 15 exact-head serializer assertions and seven native scoping component cases passed, covering configuration changes, dispatch boundaries, captures, shadowing, unused parameters, three nesting levels, and sibling filters.
cargo fmt --all --checkandgit diff --checkpassed.- The native component probes used cached DataFusion 55.0 artifacts. The five relevant lambda/HOF/filter source files are byte-identical to the exact 55.1.0 sources. The Comet cast and division implementations used also match this head, as does the relevant scalar-function factory branch.
Validation limits: the full locked native build could not resolve DataFusion 55.1.0 from the configured registry. Full exact-head JNI execution and the Comet SQL suites were not run. The second finding is supported by a native component reproduction plus exact-head Spark-oracle and serializer checks.
Current-head CI and CodeQL both report action_required, awaiting approval. Only the label workflow has passed.
| val functionsProto = expr.functions | ||
| .map { | ||
| case slf: SparkLambdaFunction => | ||
| exprToProtoInternal(slf.function, inputs, binding) |
There was a problem hiding this comment.
[P2] Preserve conditional evaluation when serializing lambda bodies
Could this conversion preserve Spark's conditional evaluation, or decline the native HOF cleanly when a speculative conversion would throw? With ANSI enabled and a Parquet column a containing [-1, 0], Spark 4.1.3 returns [] for:
SELECT filter(a, x -> CASE WHEN x > 0
THEN CAST('bad' AS INT) > 0 ELSE false END)
FROM t;Spark's normal optimizer retains the cast inside the guarded branch. The new traversal of slf.function reaches CometCast.convert, which calls cast.eval() for a literal child. That throws CAST_INVALID_INPUT during serialization even though no element takes that branch. The exception escapes before the fallback in convert can run.
I reproduced this through the exact-head public QueryPlanSerde.exprToProto entry point using the optimized Spark expression. With spark.comet.exec.higherOrderFunction.native.enabled=false, the same expression emits JVM dispatch successfully. The base implementation also dispatched the whole general filter without traversing its lambda. Please add an ANSI regression with a guarded error and input that does not take the failing branch.
There was a problem hiding this comment.
Thanks for the thorough review and the spot-on reproduction case, @sunchao!
I've addressed this by wrapping the speculative native HOF serialization in a try-catch block. When traversing the lambda body hits an expression that throws during planning (such as CometCast.convert eagerly evaluating cast.eval() in an unreachable/guarded branch under ANSI mode), we now catch NonFatal and cleanly return None. This allows convert to gracefully degrade to JVM codegen dispatch rather than letting the exception escape and abort query planning.
I've also added an ANSI mode regression test covering this exact guarded branch scenario ([-1, 0] input with filter(a, x -> CASE WHEN x > 0 THEN CAST('bad' AS INT) > 0 ELSE false END)).
| Ok(Arc::new(HigherOrderFunctionExpr::try_new_with_schema( | ||
| udf, | ||
| args, | ||
| &input_schema, | ||
| Arc::new(ConfigOptions::default()), | ||
| )?)) |
There was a problem hiding this comment.
[P2] Skip native predicate evaluation when there are no array elements
Could the native filter return the empty/null arrays before evaluating its lambda when the flattened element count is zero? With ANSI enabled, Spark safely returns [] for a nonnull empty array in partition 0:
SELECT filter(a, x -> (1 DIV spark_partition_id()) > 0)
FROM t;I verified this with normal Spark 4.1.3 optimization and Parquet input, including a null-array row. The exact-head serializer admits it as a native HOF. SparkPartitionIdBuilder lowers the partition ID to a scalar literal, so partition 0 gives the native predicate a scalar division by zero.
DataFusion's evaluate_single_list_lambda returns early for an all-null batch, but otherwise invokes the lambda even when the flattened values have length zero. A component reproduction using Comet's actual decimal_integral_div therefore raises DIVIDE_BY_ZERO on an empty array, whereas Spark never evaluates the predicate. The component dependencies and source-equivalence checks are described in the review summary; this was not a full Comet/JNI query run.
Please preserve the row null mask in the empty-elements fast path and add a regression with at least one nonnull empty array. Using a nonfoldable predicate such as this one avoids an unrelated failure from Spark constant folding.
There was a problem hiding this comment.
Thanks for the detailed analysis and the reproduction case, @sunchao!
I completely agree that evaluating the predicate on zero elements breaks Spark's short-circuit contract in ANSI mode. Looking closely at evaluate_single_list_lambda, this seems to be an upstream DataFusion gap: it already has an early return for all_null, but misses an early return when flattened.is_empty().
To avoid blocking this PR while we address it upstream, what do you think about the following approach?
Instead of a heavy physical expression wrapper around HigherOrderFunctionExpr (which would require manual array downcasting, offset inspection, and null-mask tracking), we could add a minimal guard inside our lambda wrapper (PhysicalPlanner / lambda.rs):
if batch.num_rows() == 0 {
return Ok(ColumnarValue::Array(arrow::array::new_empty_array(self.data_type())));
}There was a problem hiding this comment.
Rechecked at 5c4ba0aab36c33370dfd339eba53920621efd25b: the empty-elements issue is still present in this head. With ANSI enabled, the partition-zero predicate from the original comment still serializes natively. Spark returns [] and null, while the native component probe raises DIVIDE_BY_ZERO for both a nonnull empty array and mixed empty/null input. The all-null control returns null correctly.
The proposed zero-row guard at the lambda-body boundary looks suitable. I tested that guard in the component harness: it produces [], preserves [[], null] for mixed input, and leaves the all-null result unchanged. DataFusion can retain responsibility for rebuilding the output arrays and their null mask.
There is no lambda wrapper left in this head, so this needs an actual physical-expression adapter around body_expr before LambdaExpr::try_new. Its children() should expose the inner body and with_new_children() should rebuild the adapter so projection rewriting preserves the guard. Could you add that implementation and an ANSI regression with at least one nonnull empty array, including mixed empty/null input?
This validation used component probes plus exact-head serializer checks and Spark 4.1.3 reference execution. I have not run the full Comet/JNI query.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed 5c4ba0aab36c33370dfd339eba53920621efd25b against 58ab5f618e1e715dee06165424672fdd820cafe4.
The guarded-cast serialization fix passes the focused checks. I found one additional P2 runtime regression involving per-element AND/OR short-circuiting, described inline. The previously reported empty-elements P2 also remains unresolved. Both runtime issues should be addressed before merging.
Validation:
- Root Maven reactor
test-compilepassed with Java 17 and Spark 4.1.3. - 44 exact-head JVM assertions passed, covering serializer routing, configuration changes, dispatch boundaries, nested captures, the guarded-cast fix, and Spark reference results for both remaining findings.
- Native component probes covered nested scopes, empty/null arrays, and guarded AND/OR predicates with controls.
cargo fmt --all --checkandgit diff --checkpassed.
Validation limits: the locked native build is blocked because the configured registry cannot resolve datafusion-datasource-json 55.1.0. Full Comet/JNI execution and the Comet SQL suites were not run. Native probes used cached internal DataFusion 55.1.0 artifacts whose six relevant lambda/HOF/filter/boolean-expression source files are byte-identical to public 55.1.0, with this head's Comet division kernel and narrow harness adapters.
| let body_expr = self | ||
| .lambda_scopes | ||
| .with_scope(scope, || self.create_expr(lambda_body, body_schema))?; | ||
|
|
||
| Ok(Arc::new(LambdaExpr::try_new(param_names, body_expr)?)) |
There was a problem hiding this comment.
[P2] Preserve per-element short-circuiting in native lambda bodies
Could we preserve Spark's per-element AND/OR evaluation before admitting these lambda bodies to the native path? With ANSI enabled and a Parquet array column a containing [0, 1]:
SELECT filter(a, x -> x <> 0 AND 1 DIV x > 0) FROM t;
-- Spark: [1]
SELECT filter(a, x -> x = 0 OR 1 DIV x > 0) FROM t;
-- Spark: [0, 1]At 5c4ba0a, both expressions serialize to native HOFs, including with JVM codegen dispatch disabled. I verified that Spark 4.1.3 retains the guard on the left in the optimized expression and returns the results above. The corresponding native component probes raise DIVIDE_BY_ZERO for both expressions.
AndBuilder and OrBuilder construct DataFusion BinaryExpr. In DataFusion 55.1.0, mixed boolean batches only mask the right operand when at most 20% of rows need it. With [0, 1], division therefore runs on the zero element despite the guard. The control [0, 0, 0, 0, 1] succeeds. The base implementation dispatched the whole general filter to Spark, and the new serialization NonFatal catch cannot intercept this runtime error.
Could we preserve the evaluation mask, or fall back for predicates whose skipped branches can raise, and add ANSI regressions for both guarded AND and OR?
Validation boundary: these are native component reproductions using the current Comet division kernel and DataFusion components whose relevant source files match public 55.1.0 byte-for-byte, paired with exact-head serializer checks and Spark reference executions. Full Comet/JNI query execution remains unrun.
There was a problem hiding this comment.
Thanks for the thorough reviews and guidance, @sunchao!
Both runtime P2 issues have been resolved, and all corresponding ANSI regression tests are now passing:
1. Zero-row lambda guard (Empty & mixed arrays)
- Implemented
EmptyBatchGuardExpr, a lightweight physical expression adapter innative/core/src/execution/lambda.rs, and wrappedbody_exprbefore callingLambdaExpr::try_new. - When
batch.num_rows() == 0, it short-circuits evaluation and returnsarrow::array::new_empty_arraydirectly, avoiding scalar runtime errors (such as1 DIV spark_partition_id()). DataFusion retains responsibility for reconstructing the output array offsets and row null masks. - The adapter properly delegates
children(),with_new_children(),fmt_sql(), and satisfiesDynEq/DynHashviadyn_eq/dyn_hash.
2. Per-element short-circuiting in conditional branches (Guarded AND / OR / CASE / IF)
- Added
hasGuardedFallibleBranchandisFallibleExprinCometHighOrderFunction.scala. - Because DataFusion evaluates vectorized branches without masking when more than 20% of batch rows require evaluation (which easily triggers on small array batches like
[0, 1]), native execution cannot guarantee Spark's per-element short-circuiting in ANSI mode. - When conditional expressions (
And,Or,CaseWhen,If,Coalesce) contain potentially fallible operations (integer division, non-try casts, arithmetic overflow, out-of-bounds array indexing) in guarded branches under ANSI mode, the native path is declined and execution cleanly falls back to JVM codegen dispatch.
3. Regression SQL tests added
Added test queries under ANSI mode covering:
- Non-null empty arrays and mixed
[[], NULL]using1 DIV spark_partition_id(). - Guarded AND with
x <> 0 AND (1 DIV x) > 0on[0, 1]. - Guarded OR with
x = 0 OR (1 DIV x) > 0on[0, 1]. - Guarded CASE WHEN with
CAST('bad' AS INT)on[-1, 0].
Could you please take another look when you have time?
# Conflicts: # spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
5 participants
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?
1. Protobuf (
native/proto/src/proto/expr.proto)HigherOrderFunc,LambdaFunction, andNamedLambdaVariable.high_order_func(71) andnamed_lambda_variable(72) fields toExpr.2. Lambda Infrastructure & Scope Management (
native/core/src/execution/lambda.rs)NamedLambdaVariableby SparkexprId, preventing name shadowing or column collisions.LambdaParamsCapture(pin_unused_params) to prevent DataFusion's optimizer from pruning unused lambda parameters and preserving physical batch structure.EmptyBatchGuardExpr, a transparent physical expression adapter wrapping lambda bodies. When an input batch has 0 rows (e.g. non-null empty arrays[]or mixed[[], NULL]), it short-circuits evaluation and returns an empty array directly. This preserves Spark's ANSI contract where lambda predicates are never evaluated on empty collections (avoiding runtime errors like scalar division by zero).3. Physical Planner (
native/core/src/execution/planner.rs)PhysicalPlannerto supportHigherOrderFuncexpressions: resolves parameter field types via the HOF UDF contract, plans the lambda body under the resolved scope, and binds physicalLambdaVariableindices.4. Spark Serde & Three-Tier Execution (
CometHighOrderFunction.scala,arrays.scala)Native -> JVM Codegen -> Sparkfallback hierarchy controlled byspark.comet.exec.higherOrderFunction.native.enabledandspark.comet.exec.scalaUDF.codegen.enabled.try-catch NonFatalto gracefully decline the native path if eager evaluation occurs in unreachable branches during planning (e.g.CometCastevaluating literal arguments in guarded branches under ANSI mode).hasGuardedFallibleBranchto decline the native path when conditional operators (AND,OR,CASE WHEN,IF,COALESCE) guard fallible expressions (division, non-try casts, overflow, indexing) under ANSI mode. This accounts for DataFusion's vectorized batch short-circuit threshold (20%) and avoids runtime exceptions on elements skipped by Spark's per-element evaluation.How are these changes tested?
array_filteroperations with captures, literals, string functions, and nested structures.[[], NULL]under ANSI mode (spark_partition_id()division by zero).ANDandORpredicates on[0, 1]under ANSI mode.CASE WHENwith malformed casts on[-1, 0]under ANSI mode.EmptyBatchGuardExprshort-circuiting on empty batches.