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?

1. Protobuf (native/proto/src/proto/expr.proto)

• Added three new protobuf messages: HigherOrderFunc, LambdaFunction, and NamedLambdaVariable.
• Added high_order_func (71) and named_lambda_variable (72) fields to Expr.

2. Lambda Infrastructure & Scope Management (native/core/src/execution/lambda.rs)

• Scope Management: Introduced nested lambda variable scopes resolving NamedLambdaVariable by Spark exprId, preventing name shadowing or column collisions.
• Optimizer Anchoring: Implemented LambdaParamsCapture (pin_unused_params) to prevent DataFusion's optimizer from pruning unused lambda parameters and preserving physical batch structure.
• Empty Batch Runtime Guard: Added 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)

• Extended PhysicalPlanner to support HigherOrderFunc expressions: resolves parameter field types via the HOF UDF contract, plans the lambda body under the resolved scope, and binds physical LambdaVariable indices.

4. Spark Serde & Three-Tier Execution (CometHighOrderFunction.scala, arrays.scala)

• Three-tier execution model: Expressions follow a clean Native -> JVM Codegen -> Spark fallback hierarchy controlled by spark.comet.exec.higherOrderFunction.native.enabled and spark.comet.exec.scalaUDF.codegen.enabled.
• Safe Speculative Serialization: Wrapped lambda serialization in try-catch NonFatal to gracefully decline the native path if eager evaluation occurs in unreachable branches during planning (e.g. CometCast evaluating literal arguments in guarded branches under ANSI mode).
• ANSI Short-Circuit Guard: Implemented hasGuardedFallibleBranch to 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?

• SQL Regression Tests:
  • Standard array_filter operations with captures, literals, string functions, and nested structures.
  • Zero-row short-circuiting on non-null empty arrays and mixed [[], NULL] under ANSI mode (spark_partition_id() division by zero).
  • Guarded AND and OR predicates on [0, 1] under ANSI mode.
  • Guarded CASE WHEN with malformed casts on [-1, 0] under ANSI mode.
• Rust Unit Tests:
  • Verification of EmptyBatchGuardExpr short-circuiting on empty batches.
• Benchmarks:
  • Micro-benchmarks comparing Native Comet vs. JVM Codegen vs. vanilla Spark (showing 3.7x–8.0x speedup).

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

[sunchao]

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-compile passed 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 --check and git diff --check passed.

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.

[sunchao]

[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.

[kazantsev-maksim]

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 in native/core/src/execution/lambda.rs, and wrapped body_expr before calling LambdaExpr::try_new.
• When batch.num_rows() == 0, it short-circuits evaluation and returns arrow::array::new_empty_array directly, avoiding scalar runtime errors (such as 1 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 satisfies DynEq/DynHash via dyn_eq/dyn_hash.

2. Per-element short-circuiting in conditional branches (Guarded AND / OR / CASE / IF)

• Added hasGuardedFallibleBranch and isFallibleExpr in CometHighOrderFunction.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] using 1 DIV spark_partition_id().
• Guarded AND with x <> 0 AND (1 DIV x) > 0 on [0, 1].
• Guarded OR with x = 0 OR (1 DIV x) > 0 on [0, 1].
• Guarded CASE WHEN with CAST('bad' AS INT) on [-1, 0].

Could you please take another look when you have time?

[sunchao]

Reviewed ba319edf9. I would not approve yet: the existing short-circuit finding is only partially fixed.

[P2] Native lambdas still evaluate elements Spark skips — CometHighOrderFunction.scala:195.

I reproduced a wrong result through full Spark/Comet execution. With one Parquet row containing a=[0,1]:

SELECT filter(a, x -> x = 0 OR monotonically_increasing_id() = 0)
FROM t;

Spark and JVM dispatch return [0,1]; native execution returns [0]. Evaluating the skipped right-hand side advances the counter.

The same guard also misses:

• Guarded abs(INT_MIN): native throws with ANSI enabled.
• Guarded element_at(..., 0): native throws with ANSI disabled.
• Guarded rand(42L): native selects different elements.
• Reusing an optimized ANSI Dataset after disabling session ANSI: guarded division throws natively.

All five cases match Spark when native HOF execution is disabled. Preserve per-element evaluation masks, or conservatively dispatch affected bodies using their captured error modes.

The native build, JVM compilation, and all four SQL fixture tests passed. Additional end-to-end checks confirmed the empty-array fix, guarded-division control, captures, shadowing, and three-level nesting. Current CI still awaits approval; the full Spark-version matrix was not run.

Nothing posted to GitHub.

[kazantsev-maksim]

Thanks for the detailed follow-up and the great reproduction cases, @sunchao! The monotonically_increasing_id() and rand() examples really hit the nail on the head — speculative execution of conditionally skipped branches isn't just about arithmetic errors in ANSI mode, but can silently corrupt results for any stateful or non-deterministic expression.

I looked into maintaining a Scala-side AST whitelist/blacklist for guarded branches, but it quickly became an endless, fragile game of whack-a-mole:

• A conservative whitelist immediately broke valid, common queries like x.id IS NOT NULL AND x.name LIKE 'a%' (because Like and string predicates weren't on the list, causing fallback to JVM codegen dispatch, which cannot bind struct fields on NamedLambdaVariable).
• A blacklist is equally impractical given the sheer number of Spark expressions that can fail or hold state (element_at(0), abs(INT_MIN), date arithmetic, etc.), across varying Spark versions and captured error modes.

This points directly to your first suggestion: preserving per-element evaluation masks.

The root cause is DataFusion's BinaryExpr for AND/OR, which skips masking when the active row count exceeds its 20% threshold. In array lambdas where batches contain very few rows (e.g. [0, 1]), this threshold is almost always exceeded.

Instead of patching Scala with fragile AST inspections, what do you think about solving this on the physical execution side in Rust?
I could introduce a strict boolean physical expression (e.g., CometStrictAnd / CometStrictOr, or an adapter during lambda physical planning) that:

1. Evaluates the left child.
2. For rows that require right-hand evaluation, filters the batch (arrow::compute::filter) so the right child is evaluated only on matching rows.
3. Merges the results back.

This would:

• Guarantee Spark's strict per-element short-circuiting semantics for all expressions (DIV, monotonically_increasing_id, element_at, abs, etc.).
• Completely eliminate the need for AST whitelists/blacklists and captured error mode checks in Scala.
• Keep all valid predicates (including LIKE, nested structs, etc.) running on the fast native path.

Does implementing this strict masking adapter in Rust for lambda bodies sound like the right direction to you, or did you have a different mechanism in mind for preserving the evaluation mask?

[sunchao]

Yes, strict masking in the native lambda path is the direction I had in mind. It addresses the evaluation semantics directly and keeps supported predicates native without maintaining a list of potentially dangerous expressions.

A few details to preserve:

• For AND, evaluate the RHS when the LHS is true or null. For OR, evaluate it when the LHS is false or null. This matters for SQL’s three-valued logic.
• Skip RHS evaluation entirely when no elements need it, including zero-row inputs, and preserve element order so stateful expressions advance correctly.
• Apply this recursively within lambda bodies and preserve it through expression rewrites.

DataFusion’s PhysicalExpr::evaluate_selection looks worth reusing for filtering and scattering, with explicit handling for empty input.

Please keep the existing empty-batch guard and speculative-serialization fallback. Runtime masking does not prevent exceptions during serialization. Before removing the Scala conditional guard entirely, we should also verify IF, CASE WHEN, and COALESCE.

Could you add native-path regressions for the five reported cases, nullable boolean conditions, and nested predicates? Then rerun the benchmarks to quantify the masking overhead. That would give us a solid basis for reviewing the implementation.

[kazantsev-maksim]

Thanks for the guidance, @sunchao! Moving strict masking into the native physical path via ShortCircuitBinaryExpr cleanly resolves the evaluation semantics without having to maintain fragile AST whitelists/blacklists in Scala.

Implementation Details

1. ShortCircuitBinaryExpr (lambda.rs):

   • SQL Three-Valued Logic (3VL):
     • AND: RHS is evaluated only when LHS is TRUE or NULL (skipped when LHS is strictly FALSE).
     • OR: RHS is evaluated only when LHS is FALSE or NULL (skipped when LHS is strictly TRUE).
   • Zero-row / unneeded RHS short-circuit: When no elements require RHS evaluation (true_count == 0), RHS evaluation is skipped entirely. This protects stateful expressions (monotonically_increasing_id, rand) and fallible operations (1 DIV x, element_at(..., 0), abs(INT_MIN)).
   • Masked evaluation via evaluate_selection: When a subset of rows requires the RHS, we evaluate via DataFusion's PhysicalExpr::evaluate_selection, which filters the batch and scatters the result back into place while preserving original element order. The arrays are then combined using Arrow's Kleene boolean logic (and_kleene / or_kleene).
   • Implements children(), with_new_children(), and fmt_sql() to seamlessly support optimizer projection rewriting.

2. Recursive Planning (planner.rs):

   • Added rewrite_short_circuit_binary in PhysicalPlanner, which recursively walks the planned lambda body tree and replaces BinaryExpr (Operator::And and Operator::Or) with ShortCircuitBinaryExpr.
   • Preserved EmptyBatchGuardExpr to protect zero-row batches at the outer lambda boundary.

3. Scala Serde Cleanup & Fallbacks (CometHighOrderFunction.scala):

   • Removed AST whitelists for AND / OR. Supported expressions such as LIKE, string functions, and nested struct access (x.id IS NOT NULL AND x.name LIKE 'a%') now remain fully native.
   • Retained try-catch NonFatal during speculative serialization to prevent eager planning exceptions (e.g. cast.eval() in guarded branches).
   • Retained codegen dispatch fallback for guarded fallible branches in CASE WHEN, IF, and COALESCE.

4. Documentation Updates:

   • Updated the ## lambda_funcs documentation table: filter is now documented as Hybrid (single-argument lambdas and array_compact run natively with strict per-element masking; multi-argument lambdas with index and unsupported shapes fall back to JVM codegen dispatch).
   • Updated ScalaDoc on CometHighOrderFunction and configuration docstrings.

Regression Test Coverage

Added native SQL fixture tests under ANSI mode covering:

• All 5 reported cases:
  • Stateful counter: x = 0 OR monotonically_increasing_id() = 0 (counter no longer advances speculatively).
  • PRNG state: x = 0 OR rand(42L) > 0.5.
  • Non-ANSI runtime failure: x = 0 OR element_at(array(1, 2), 0) = 1 (index 0 is protected).
  • Arithmetic overflow: x = 0 OR abs(-2147483648) > 0.
  • Division by zero: x <> 0 AND (1 DIV x) > 0 and x = 0 OR (1 DIV x) > 0.
• Nullable boolean conditions (3VL):
  • Evaluates RHS when LHS is NULL (e.g. (x > 0) OR (x IS NULL) preserving NULL elements).
• Nested predicates:
  • Multi-level compound predicates combining AND and OR with division guards.
• Empty & mixed array batches:
  • [] and NULL rows with 1 DIV spark_partition_id().
• Baseline control:
  • Standard comparisons (x > 0 AND x < 10) continue to execute natively.

[sunchao]

Summary

• Prior state and problem: General filter lambdas used JVM dispatch. This PR moves supported predicates into DataFusion to avoid that execution overhead.
• Design approach: Adds HOF protobuf messages, expression-ID-based lambda scopes, native planning, and native → JVM → Spark fallback.
• Correctness / compatibility analysis: The existing P2 short-circuit report remains unresolved. For Parquet input a=[0,1], filter(a, x -> x = 0 OR monotonically_increasing_id() = 0) returns [0,1] in Spark 4.1.3 but [0] in the native component reproduction. CometHighOrderFunction.scala:195 still admits this predicate, and DataFusion evaluates the skipped RHS, advancing its counter. The strict masking described in the discussion is absent from this head. Spark's relevant ArrayFilter semantics agree across supported 3.4–4.2 sources.
• Key design decisions: Reusing upstream lambda projection and array filtering limits custom machinery. The expression-ID scope stack handles binding, and EmptyBatchGuardExpr preserves empty-input behavior. The conditional fallibility check remains insufficient for Spark semantics.
• Implementation sketch: Reviewed all 15 changed files, including serde registration, configuration, native planner, protobuf additions, SQL fixtures, and benchmarks.
• Behavioral changes worth calling out: Native unary filters are enabled by default. Indexed lambdas retain dispatch fallback. Empty, mixed empty/null, and all-null component checks passed. Benchmarks cover captures, strings, nulls, and nested arrays, but were not rerun.
• Suggested improvements: Resolve the existing blocker with strict per-element evaluation masks or conservative dispatch, and assert native execution for its regression tests. No additional introduced P1/P2 issues were found within this review.

Reviewed full SHA 670847d7a3ca27e32d0eb23e1f8f5d26c60fce03 against 466e3fe3a7200f35113b523249352b5a33a21365. Routed skills: review-comet-pr and review-comet-expression-pr. Read existing discussion and threads, excluding Copilot reviews.

Exact-head CI: Comet CI and CodeQL report action_required. Only labeling passed.

Validation limits: Native probes compiled this head's guard and counter against cached DataFusion 55.1.0 components, with relevant sources verified byte-identical to upstream. Spark reference execution passed with ANSI both enabled and disabled. Full Comet/JNI execution and SQL suites were not run: Maven bootstrap failed DNS resolution, the system Maven rejected repository configuration syntax, and the native build lacked JNI headers. The Spark-version test matrix was not run. git diff --check passed.

[sunchao]

Summary

• Prior state and problem: General filter lambdas used JVM dispatch. This PR introduces native evaluation to reduce that overhead.
• Design approach: Adds HOF protobuf messages, expression-ID-based scopes, DataFusion lambda planning, and native → JVM → Spark fallback.
• Correctness / compatibility analysis: The existing P2 short-circuit concern remains unresolved. For Parquet input a=[0,1], filter(a, x -> x = 0 OR monotonically_increasing_id() = 0) returns [0,1] in Spark 4.1.3 with ANSI enabled or disabled. The native component reproduction returns [0] because evaluating the skipped branch advances the counter. The strict masking described in the discussion is absent from this head. Spark's relevant ArrayFilter semantics agree across supported 3.4–4.2 sources.
• Key design decisions: Reusing DataFusion's lambda projection and filtering limits custom machinery. Expression-ID scopes handle binding, and EmptyBatchGuardExpr correctly preserves empty-input behavior. The conditional fallibility check does not cover the existing blocker.
• Implementation sketch: Reviewed the full base-relative diff across all 15 changed files, surrounding code, existing reviews and discussion, excluding Copilot.
• Behavioral changes worth calling out: Native unary filters are enabled by default. Indexed lambdas retain dispatch fallback. Empty, mixed empty/null, and all-null component checks passed. Benchmarks cover captures, strings, nulls and nested arrays, but were not rerun.
• Suggested improvements: Resolve the existing blocker with strict per-element evaluation masks or conservative dispatch, and assert native execution in its regression tests. No additional introduced P1/P2 issues found within this review.

Reviewed full SHA 670847d7a3ca27e32d0eb23e1f8f5d26c60fce03 against 466e3fe3a7200f35113b523249352b5a33a21365. Routed skills: review-comet-pr and review-comet-expression-pr.

Exact-head CI: Comet CI and CodeQL report action_required. Only labeling passed.

Validation: Recompiled this head's guard and counter in the native component harness against cached DataFusion 55.1.0 artifacts. Verified relevant dependency sources against upstream. Reran Spark 4.1.3 reference execution and git diff --check successfully. Full locked native build, Comet/JNI SQL suites, benchmarks and the Spark-version execution matrix were not rerun. Prior exact-head build logs record Maven bootstrap/configuration and JNI-header failures, so full integration remains unverified.

[kazantsev-maksim]

Thanks for catching that, @sunchao. You were completely right — the ShortCircuitBinaryExpr implementation was accidentally missing from the previous commit. I have now properly committed and pushed all the Rust and physical planning changes to the branch.

To address your point about asserting native execution and prevent any silent JVM codegen fallback, I have also explicitly disabled the codegen dispatcher (spark.comet.exec.scalaUDF.codegen.enabled=false) in the regression tests where native execution is expected. This guarantees that DataFusion's native runtime is directly exercised.

All tests now pass natively without fallback, including:

• Stateful expressions (monotonically_increasing_id() counter no longer advances when skipped, returning [0, 1])
• PRNG sequence preservation with rand()
• element_at(..., 0) and abs(INT_MIN) in guarded branches
• Guarded division on [0, 1] under ANSI mode
• SQL 3VL nullable boolean conditions and nested predicates

Could you please take another look when you have a chance?

[sunchao]

Re-reviewed 8f5bc62341e3f4394e1bd35b87ca2ac15f977288 against 605051ad239ef704f5f25d67910a446a6b6d7c70 with five independent scopes covering native planning/binding, Spark serde and fallback, runtime errors and memory, compatibility, and performance/tests. Findings were independently reproduced and challenged.

The strict boolean adapter fixes the earlier AND/OR examples: guarded abs(INT_MIN), legacy element_at(..., 0), seeded rand, monotonically_increasing_id, and a reused ANSI Dataset after changing the session setting now agree with Spark. Empty/mixed-null arrays and nested capture/shadowing controls also pass. Reusing DataFusion's lambda infrastructure and keeping the masking adapter local is a reasonable design. Indexed/unsupported lambdas retain JVM dispatch when it is enabled; I found no additional unintended Spark row fallback.

Three P2 issues remain, detailed inline: CASE metadata checks consume state, COALESCE duplicates nondeterministic child evaluation, and captured arrays can cause multiplicative allocation growth. The two correctness issues have distinct causes and both reproduce through the complete current-head native execution path with JVM dispatch disabled. I am requesting changes for these issues.

Validation:

• Built native code with cargo build --locked -p datafusion-comet --no-default-features, then ran Maven from the root reactor. All four CometSqlFileTestSuite array_filter fixtures passed (4 tests, 1 suite, no failures).
• Ran fresh Spark 4.1.3 / JVM-dispatch / native comparisons on Parquet inputs, including the two wrong-result cases, controls, the earlier reported failures, and captured-array routing. Native probes required CometProject with dispatch disabled.
• Focused native checks covered nested binding, scalar/array three-valued boolean masks, empty/null inputs, and capture allocation growth.

Current-head CI: Comet CI and CodeQL remain action_required. Only the label check has passed.

Limits: full upstream Spark SQL and the Spark-version execution matrix were not run. Release benchmarks were not rerun. Allocation numbers are from a focused native component measurement, not a complete-query memory or timing benchmark.

[sunchao]

[P2] Keep lambda nullability checks from consuming predicate state

Could the lambda wrapper report nullability without evaluating stateful predicates? With one Parquet row containing a=[1,2,3], this query returns [1] in Spark and JVM dispatch, but [] through the native path at this head:

SELECT filter(a, x ->
  (CASE WHEN monotonically_increasing_id() = 0 THEN x ELSE 0 END) > 0)
FROM t;

HigherOrderFunctionExpr asks the lambda body for its return field during construction and evaluation. This delegation reaches DataFusion's CaseExpr::nullable(), which evaluates the WHEN condition on a synthetic one-row batch and advances the same counter later used for real elements. The conditional guard admits the expression because its result branches are not fallible. The previous whole-filter JVM dispatcher does not perform these evaluations. Please keep metadata inspection free of these side effects and add a native regression with a stateful CASE condition. I reproduced this through full Spark 4.1.3/Comet execution with JVM dispatch disabled.

[sunchao]

[P2] Preserve single evaluation of nondeterministic COALESCE children

Could native admission also account for nondeterministic nonfinal COALESCE arguments? With one Parquet row a=[1,2,3,4], this returns [2,4] in Spark and JVM dispatch, but [4] with native HOF execution and dispatch disabled:

SELECT filter(a, x ->
  coalesce(IF(rand(42L) < 0.6, x, CAST(NULL AS INT)), -1) = x)
FROM t;

CometCoalesce serializes the child separately as WHEN IS NOT NULL(child) and THEN child. The second random-expression instance runs only on the selected subset, so its draws correspond to different elements. This remains wrong even when CASE metadata evaluation is suppressed, so it is separate from the nullability issue. Plain IF with the same seeded predicate and deterministic COALESCE controls both agree with Spark. Please ensure each child is evaluated once per element, or retain whole-filter JVM dispatch for these shapes until the native lowering preserves that contract. This mismatch is reproduced through full current-head Spark 4.1.3/Comet execution.

[sunchao]

[P2] Avoid copying captured arrays once per lambda element

Could we prevent deep replication of captured arrays before selecting this path by default? For filter(a, x -> x >= 0 AND size(b) > 0), DataFusion's evaluate_single_list_lambda broadcasts captured b with take_arrays, copying its entire child buffer once for each element of a. With one row containing 4,096 integers in each array, a focused native allocation probe measured 67,705,058 bytes allocated from 33,200 bytes of input. The repeated integer payload alone is 64 MiB, although this predicate only needs b's length. The cost grows with both array lengths and batch row count.

Spark retains the outer-row array reference, and the base implementation dispatches this filter to Spark's evaluator. I confirmed that Catalyst keeps size(b) inside the lambda and that the current Comet query takes the native path with dispatch disabled. Please avoid expanding complex captures this way, or retain the existing JVM dispatcher for affected shapes until native captures can represent them efficiently. An allocation benchmark with a captured array would cover this regression; the current capture benchmark uses a scalar integer. These allocation numbers are component measurements, not complete-query memory or timing results.

[sunchao]

Summary

• Prior state and problem: General filter lambdas used JVM dispatch. This PR introduces native evaluation to reduce that overhead.
• Design approach: Adds HOF protobuf messages, expression-ID-based lambda scopes, DataFusion planning, and native → JVM → Spark fallback.
• Correctness / compatibility analysis: Found one new P2 regression: null-intolerant comparisons evaluate operands that Spark skips, changing stateful predicate results. Spark’s relevant semantics agree across supported 3.4–4.2 sources. The existing CASE and COALESCE wrong-result reports also reproduce at this head.
• Key design decisions: Reusing DataFusion’s lambda machinery limits custom code. The strict boolean adapter correctly handles tested three-valued logic and masked evaluation, but its scope does not cover implicit null short-circuiting. The existing captured-array replication concern remains substantiated by the take_arrays capture path.
• Implementation sketch: Reviewed the entire base-relative diff across all 16 files, including surrounding planner and serde code, configuration, protobuf, SQL fixtures, benchmarks, and documentation.
• Behavioral changes worth calling out: Supported unary filters execute natively by default. Indexed and unsupported shapes retain fallback. Nested capture, shadowing, empty-input, and boolean-mask controls passed.
• Suggested improvements: Preserve null-dependent operand evaluation and add a native regression for the new finding. Resolve the three blockers in the existing review, including captured-array allocation growth.

Reviewed full SHA 8f5bc62341e3f4394e1bd35b87ca2ac15f977288 against 605051ad239ef704f5f25d67910a446a6b6d7c70. Routed skills: review-comet-pr and review-comet-expression-pr. Read existing reviews, discussion, inline comments, and threads, excluding Copilot.

Exact-head CI: Comet CI and CodeQL remain action_required. Only labeling passed.

Validation: Locked native build with --no-default-features and root Maven compilation passed. All four array_filter SQL fixtures passed. Ran 36 Spark/JVM-dispatch/native comparisons with ANSI enabled and disabled. Native comparisons required CometProject with dispatch disabled. Focused serializer and native boolean-mask checks also passed.

Validation limits: Full upstream Spark SQL, the Spark-version execution matrix, lint, and release benchmarks were not run. The existing capture-allocation measurement was not rerun.

[sunchao]

[P2] Preserve null short-circuiting for comparison operands. With one Parquet row containing a=[NULL,0] as ARRAY<BIGINT>, filter(a, x -> x = monotonically_increasing_id()) returns [0] in Spark and JVM dispatch but [] natively. Spark’s BinaryExpression.eval skips the RHS when the LHS is null. This rewrite leaves equality as DataFusion BinaryExpr, which evaluates the counter for both elements, so the comparison against zero sees counter value 1. The newly enabled native path therefore silently drops a valid element. Could it preserve the null-dependent evaluation mask for affected expressions, or dispatch those lambda bodies until that contract is supported?

Evidence: Reproduced through full exact-head Comet execution on Spark 4.1.3 in local[1]. Write SELECT CAST(array(NULL, 0) AS ARRAY<BIGINT>) AS a to Parquet, register it as t, then run SELECT filter(a, x -> x = monotonically_increasing_id()) FROM t. With spark.comet.exec.higherOrderFunction.native.enabled=true and spark.comet.exec.scalaUDF.codegen.enabled=false, the asserted CometProject returns [], with ANSI either enabled or disabled. Spark and whole-filter JVM dispatch both return [0]. Adding x IS NOT NULL AND restores [0] natively. A component reproduction using this head’s lambda adapter and counter independently produced the same mismatch.

[kazantsev-maksim]

Thanks @sunchao for the thorough review and for pointing out the memory replication in take_arrays as well as the boolean conditional discrepancies, @sunchao!

I have addressed both blockers by refining the degradation rules in CometHighOrderFunction.scala:

1. Complex Outer Captures (capturesComplexOuterAttribute)

• Lambda bodies that capture outer columns of complex types (ArrayType, MapType, StructType) now safely degrade to JVM codegen dispatch.
• This avoids quadratic memory replication from DataFusion's take_arrays (such as the 64 MiB allocation on 4k elements).
• Capturing scalar columns (Int, String, etc.) remains fully native and fast.

2. Boolean Conditional Expressions (hasUnsupportedConditionals)

• Conditional expressions (CaseWhen, If, Coalesce) that return BooleanType and serve as filter predicates now degrade to JVM codegen dispatch. This avoids DataFusion's native evaluation discrepancies ([1] -> [] and [2, 4] -> [4]).
• Scalar conditionals (e.g. coalesce(x, 0) > 0) continue to execute natively in DataFusion.

3. Tests & Documentation

• Added regression tests covering:
  • Complex outer array capture: filter(a, x -> x >= 0 AND size(b) > 0)
  • Boolean COALESCE: filter(a, x -> coalesce(x % 2 = 0, false))
  • Boolean CASE WHEN: filter(a, x -> CASE WHEN x = 1 THEN true ELSE x > 2 END)
  • Boolean IF: filter(a, x -> if(x = 2, false, true))
  • Verified that scalar conditionals (filter(arr, x -> coalesce(x, 0) > 0)) continue to execute natively with codegen disabled.
• Updated the ScalaDoc on CometHighOrderFunction and the corresponding PR documentation.

All native and fallback tests pass. Could you please take another look when you have a chance?

[andygrove]

This is a light fully automated review since there are so many PRs open.

docs/source/contributor-guide/roadmap.md:48-50 still describes DataFusion's array_filter as native higher-order function support "that Comet does not yet use". Once this lands with spark.comet.exec.higherOrderFunction.native.enabled defaulting to true, unary filter lambdas run through that function by default, so the "Native Lambda Evaluation" section becomes wrong. Could that paragraph be updated in this PR? A few code comments also describe code that no longer exists. Item 3 of the module doc at native/core/src/execution/lambda.rs:28-30 describes a wrapper that keeps unused lambda parameters visible in children(), which went away when DataFusion 55's LambdaExpr started computing used_param_indices() itself. The module doc also never mentions EmptyBatchGuardExpr or ShortCircuitBinaryExpr, which are what the file contains now. native/core/src/execution/planner.rs:3850 says "the guard pops on any ? / drop", but with_scope pops explicitly and there is no guard anymore. The Scaladoc at spark/src/main/scala/org/apache/comet/serde/CometHighOrderFunction.scala:169 points at StrictBooleanExpr, which should be ShortCircuitBinaryExpr.

[sunchao]

Summary

• Prior state and problem: General filter lambdas used JVM dispatch. This PR introduces native evaluation to reduce that overhead.
• Design approach: Adds HOF protobuf messages, expression-ID-based lambda scopes, DataFusion planning, and native → JVM → Spark fallback.
• Correctness / compatibility analysis: Reproduced the existing CASE ([1] → []), COALESCE ([2,4] → [4]), and null-comparison ([0] → []) regressions through full native execution, with ANSI enabled and disabled. Spark and JVM dispatch agree. Relevant Spark semantics match across supported 3.4–4.2 sources.
• Key design decisions: Reusing DataFusion’s lambda infrastructure limits custom code. The strict boolean adapter passes focused masking and three-valued-logic checks. The existing captured-array allocation concern remains substantiated by the take_arrays replication path.
• Implementation sketch: Reviewed the entire 16-file base-relative diff and surrounding code, including serialization, binding, configuration, SQL fixtures, benchmarks, and documentation.
• Behavioral changes worth calling out: Supported unary filters run natively by default. Indexed and unsupported shapes retain fallback. Nested captures, shadowing, empty inputs, and boolean-mask controls passed.
• Suggested improvements: Resolve the existing CASE, COALESCE, and capture-allocation blockers and null short-circuiting blocker. Preserve side-effect-free metadata inspection, single evaluation, operand masks, and bounded capture allocation. No additional introduced P1/P2 issues found within this review.

Reviewed full SHA 8f5bc62341e3f4394e1bd35b87ca2ac15f977288 against 605051ad239ef704f5f25d67910a446a6b6d7c70. PR remains non-draft. Routed skills: review-comet-pr, review-comet-expression-pr, and review-comet-memory-pr. Read existing discussions and review threads, excluding Copilot.

Exact-head CI: Comet CI and CodeQL remain action_required. Only labeling passed.

Validation: Locked native build with --no-default-features passed. Root Maven reactor ran all four array_filter fixtures successfully. Completed 36 Spark/JVM-dispatch/native executions with native comparisons requiring CometProject and dispatch disabled. Focused Rust boolean-mask checks passed.

Validation limits: Full upstream Spark SQL, the Spark-version execution matrix, lint, release benchmarks, and the prior allocation measurement were not rerun. Project source remains unchanged.