fix: match Spark's null short-circuiting in array_join and enable it natively

[Visorgood]

Which issue does this PR close?

Closes #3178.

Rationale for this change

array_join was flagged Incompatible with the note "Null handling may differ from Spark", but the specific difference was never pinned down. It turned out to be two things, neither of them what #3178 guessed at.

1. A null nullReplacement does not nullify the row. Spark returns null whenever nullReplacement evaluates to null, even when the array holds no nulls to replace. DataFusion's array_to_string reads a null null_string as "omit null elements", i.e. the same as not passing a third argument:

arr	nullrep	Spark	Comet (before)
["a", null, "c"]	NULL	NULL	a,c
["a", "b", "c"]	NULL	NULL	a,b,c

2. Spark short-circuits past later arguments; DataFusion does not. ScalarFunctionExpr evaluates every argument up front, so an argument that can throw fails on rows where Spark returns null: array_join(arr, element_at(delims, 0)) with arr = NULL throws instead of returning null. Worse, ArrayJoin.eval and ArrayJoin.doGenCode disagree on the order – eval checks the array first, doGenCode evaluates the replacement first – so there is no single ordering a native lowering could reproduce for both. Thanks to @sunchao for finding this and for suggesting the resolution below.

The two cases #3178 actually asks about – null elements skipped without a replacement, and substituted with one – were already correct; the tests here lock that in.

What changes are included in this PR?

Only order-insensitive arguments run natively. getSupportLevel reports Compatible when the delimiter and the null replacement are each a literal or a column read. Those cannot throw, carry state or have side effects, so evaluating them earlier than Spark would is unobservable and the ordering question does not arise. Everything else reports Incompatible and runs through the JVM codegen dispatcher, which is doGenCode itself and therefore correct by construction. Note that foldable is not usable as this test: ConstantFolding deliberately leaves a throwing foldable expression unfolded inside a conditional branch.

One null guard. convert wraps the array_to_string call in IsNull(nullReplacement) when the replacement is nullable, which is the fix for (1). The array argument is unrestricted, since it is evaluated on every path and the other two cannot throw.

Support level. Non-collated array_join now runs natively by default instead of always going through the dispatcher. Non-default string collations remain Incompatible under #2190, mirroring CometReverse. This is strictly more native than main, where array_join never runs natively at all, and it reduces dispatcher usage rather than increasing it.

Docs. The audit entry and the expressions.md note describe the new rule.

No native or protobuf changes were required.

How are these changes tested?

Comet SQL Tests under spark/src/test/resources/sql-tests/expressions/array/:

• array_join.sql: expanded from 2 queries to 15: null elements leading / trailing / only / all, empty-string elements versus nulls, null delimiters, empty and multi-character delimiters, non-string element types, an empty-string replacement, and every combination of column and literal arguments.
• array_join_null_replacement.sql: regression coverage for (1), including the array-with-no-nulls case. The replacement is a column so these take the guarded native path.
• array_join_eager_eval_dispatch.sql: pins the dispatch routing for (2), using element_at(delims, 0) as an argument that throws when evaluated.

CometArrayExpressionSuite:

• the existing array_join test no longer forces allowIncompatible, so the default path is exercised over the 10000-row dictionaryEnabled Parquet matrix, including the guarded shape and a literal NULL replacement. That last one only reaches the native path with constant folding enabled, which the sql-tests suite disables, so it can only be covered here.
• array_join support level pins the native path and array_join guards only a nullable replacement assert the verdict and the guard placement directly. Because CometArrayJoin is a CodegenDispatchFallback, an Incompatible verdict runs Spark's own doGenCode and produces identical results, so every result assertion would keep passing if the support level silently regressed. Each of these was confirmed to fail when the behaviour it pins is removed.

Verified locally with CometArrayExpressionSuite in full alongside the sql-tests, scalastyle and spotless enabled:

Spark	Scala	JDK	Result
4.1.3	2.13	21	63 passed
4.0.4	2.13	17	63 passed
3.5.9	2.12	21	63 passed
3.4.3	2.12	17	59 passed, 4 cancelled (isSpark35Plus gates)

apache-rat and scalafix (syntactic and semantic) are also clean.

One shape has no SQL fixture: array_join(arr, element_at(array(','), 0)), where the
delimiter is foldable but still throws. It fails to compile the dispatcher's generated Java
("project_isNull_2" is not an rvalue) while plain Spark answers it correctly, and it still
reproduces after #5623, which fixed the adjacent foldable-subtree case in
CometBatchKernelCodegen.canShortCircuitNulls – this one is a compile failure in Spark's
whole-stage codegen rather than a bad short-circuit, so it is a separate pre-existing bug and
not something introduced here. Its verdict is pinned in CometArrayExpressionSuite instead,
which does not execute the query.

[andygrove]

Thanks for digging into this. The diagnosis is solid. I read through Spark's ArrayJoin.eval and DataFusion's generate_string_array and both match your description, including the part where a null replacement nullifies the result even when the array holds no nulls to replace.

A few things I checked while reviewing, in case they save anyone else the trip. The IfExpr guard shape is consistent with how CometArraysZip already handles its own null propagation in the same file, so the construction is idiomatic. Dictionary-encoded columns are unpacked in ScanExec::get_next, so the extra null_strings argument cannot hit array_to_string's "unsupported type for third argument" arm. The synthetic IsNull node is not part of op.expressions, so rollUpInfoMessages will not pollute the extended-explain coverage counts with an IsNull the user never wrote. And the Current status: audit bullet matches the format array_intersect already uses right above it.

Four things I would like to see addressed before this merges.

CI has not run. gh pr checks reports nothing at all and the status rollup is empty, so a maintainer will need to approve the workflows. Since this moves array_join onto the native path by default on every supported Spark version, I would want the 3.4 and 3.5 jobs green and not just the local runs in the description.

The Scala test still forces the opt-in. CometArrayExpressionSuite.scala:522 wraps the whole array_join test in withSQLConf(CometConf.getExprAllowIncompatConfigKey(classOf[ArrayJoin]) -> "true"). That was correct before this change and is now stale. It matters more than it looks, because that test is the only array_join coverage that runs over real Parquet with the dictionaryEnabled matrix and 10000 rows. With the wrapper still in place, the default path never gets exercised on that input shape. Could it be dropped?

The new tests cannot distinguish native from the dispatcher. CometArrayJoin is a CodegenDispatchFallback, so an Incompatible verdict routes through Spark's own doGenCode and produces Spark-identical results. checkSparkAnswerAndOperator checks operator coverage rather than which path an individual expression took. If getSupportLevel ever went back to Incompatible, every query in array_join_null_replacement.sql would still pass, silently. That is the same vacuous-pass shape the suite already guards against for expect_error in requireSentinelForCodegenExpectError, and it is the mirror image of the point you make in the description about the old fixture. Would it be worth adding something that pins the native path explicitly, perhaps in CometArrayExpressionSuite?

Two test cases worth adding to array_join.sql. Spark's inputTypes accepts any array that implicitly casts to array<string>, so array_join(array(1, 2, 3), ',') and array_join(array(1, NULL, 3), ',', 'X') are both valid and fairly common in practice. Neither file covers a non-string element type today. It might also be worth adding an empty-string replacement such as array_join(array('a', NULL, 'b'), ',', ''). The file covers empty-string elements well, but '' versus NULL for the replacement is exactly the distinction the new guard draws, so it seems worth pinning.

[sunchao]

[P2] Keep delimiter evaluation behind the null-array guard

Making this Compatible changes the default behavior of array_join(arr, element_at(delims, idx)) over supported Parquet input columns. For a row with arr = NULL, delims = [','], and idx = 0 (arr/delims are ARRAY<STRING> and idx is INT), Spark's generated code returns NULL without evaluating the delimiter. The native ScalarFunctionExpr evaluates every argument first, so ListExtract throws for index zero before array_to_string can inspect the null array, even with ANSI disabled. The previous default JVM dispatcher preserved Spark's guard, and the native error propagates rather than falling back. Please preserve that guard on this newly enabled path. This is source reasoning for the inspected Spark 3.5/4.0 paths; I have not run it.

[sunchao]

[P2] Follow-up on the remaining null-array case

Thanks, the original column-based delimiter example is guarded now. Could the guard selection also distinguish literal values from surviving foldable expressions? With Parquet columns flag = true and arr = NULL, consider SELECT IF(flag, array_join(arr, element_at(array(','), 0)), NULL) FROM t. Spark's ConstantFolding deliberately retains a foldable expression that throws inside a conditional branch, so this does not require disabling the optimizer. The delimiter remains foldable, which makes the new afterArray.exists(!_.foldable) test omit the array guard. Spark's two-argument generated path returns NULL without evaluating it, while the native scalar-function path evaluates ListExtract and raises INVALID_INDEX_OF_ZERO.

Could the guard also protect a surviving foldable delimiter, with a regression using normal constant folding? This is the remaining part of this null-array issue, based on the maintained Spark 3.5/4.0 sources and the current native path. I have not executed the query.

[sunchao]

[P2] Evaluate a nullable replacement once per input row

The replacement is serialized both as the function argument and again under IsNull. Consider IF(monotonically_increasing_id() % 2L = 0L, CAST(NULL AS STRING), 'X') with a sole input array column containing four rows of ['a', NULL, 'b'] in one native batch. The two native replacement trees have separate counters: the guard sees all four rows, while the function's copy sees only the two rows that passed the guard. That second copy produces one NULL replacement, which array_to_string treats as omission, yielding 'a,b'. Spark's single evaluation can produce only NULL or 'a,X,b'. The mismatch survives either starting parity, including counters advanced during planning. Please compute the replacement once and reuse its value for both the guard and join. This is source reasoning for the inspected Spark 3.5/4.0 paths, not an executed reproduction.

[Visorgood]

Thanks @andygrove and @sunchao! These were good catches. Pushed a fix that addresses them.

eager argument evaluation
Confirmed and fixed. I reproduced it before changing anything: with no guard, SELECT array_join(arr, element_at(delims, idx)) FROM t WHERE arr IS NULL throws inside Comet rather than returning null, exactly as was described.

convert now nests the array_to_string call inside IsNull guards in Spark's own evaluation order (array, then delimiter, then replacement). This works because Comet's IfExpr delegates to DataFusion's CaseExpr, which evaluates each branch against a filtered remainder_batch – so a guarded argument is never evaluated on the short-circuited rows. New fixture array_join_null_array_guard.sql covers it; removing the guard makes it fail.

Guards are emitted only where they protect something: earlier arguments are guarded only when a later one is non-foldable, and non-nullable arguments never are. So array_join(arr, ',') and array_join(arr, ',', 'X') keep the exact plan they had before this PR.

double evaluation of the replacement
Indeed there is no way to bind a value once in a serialized expression tree. Rather than paper over it, getSupportLevel now reports Incompatible when any argument that needs a guard is non-deterministic, so those route through the codegen dispatcher where Spark's single evaluation is preserved. The monotonically_increasing_id() % 2 example takes that path.

stale withSQLConf
Dropped. That test now exercises the default path over the 10000-row dictionaryEnabled Parquet matrix, plus one added query with a nullable non-foldable delimiter and replacement to cover the guarded shape.

tests cannot distinguish native from the dispatcher
Agreed. I checked, and reverting Compatible() to Incompatible left every .sql fixture green. Added array_join support level pins the native path, which asserts the verdict directly, and array_join emits a null guard only where it is needed, which asserts the guard is emitted only in the cases above. I verified both fail when the behavior they pin is removed.

extra test cases
Added non-string element types (array(1, 2, 3), array(1, NULL, 3) with and without a replacement, plus decimal and boolean) and the empty-string replacement, since '' versus NULL is exactly the distinction the guard draws.

Docs: the audit entry and the expressions.md note now describe the guards and the non-deterministic carve-out.

Verified locally on Spark 3.4.3, 3.5.9, 4.0.4 and 4.1.3 with CometArrayExpressionSuite in full alongside the sql-tests, scalastyle and spotless enabled: 52 passing on 4.0/3.5/4.1, 50 plus 2 isSpark35Plus cancellations on 3.4.

On CI: the earlier macos-14/Spark 4.0 [scans] failure looks unrelated to this change. Could someone re-approve the workflows so this gets a clean run?

[sunchao]

Re-reviewed b5f6d3effd2dfb22ce018c10058b8829557c219b. The original column-delimiter case and nullable nondeterministic replacement case are addressed. I found a new P2 in generated-code evaluation order and a remaining foldable-delimiter case for the existing null-array discussion.

These are source-derived findings. I did not execute Spark/Comet queries or tests. CI currently includes a failing Spark 3.5 Build Native + JVM Test Classes job, whose cause I have not classified.

[sunchao]

[P2] Preserve the generated code's replacement-first evaluation

Could these guards preserve ArrayJoin.doGenCode's evaluation order? In the maintained Spark 3.5/4.0 implementations, the replacement is evaluated and null-checked before the array and delimiter. With nullable Parquet columns arr = ['a','b'], delims = [','], idx = 0, and nullrep = NULL, array_join(arr, element_at(delims, idx), nullrep) therefore returns NULL in Spark's generated path. Here the deterministic arguments remain Compatible, and the delimiter guard runs before the replacement guard, so ListExtract raises INVALID_INDEX_OF_ZERO, even with ANSI disabled. The previous revision's outer replacement guard skipped that delimiter.

This also matters for a non-nullable replacement such as CAST(monotonically_increasing_id() AS STRING). It is excluded from guardedArgs, so the nondeterminism check permits it, but the new array guard advances its counter only for retained rows. Spark evaluates it before checking the array on every row. Moving only the nullable replacement guard would leave this case.

Could the lowering preserve replacement-first, single evaluation, or use the dispatcher for cases it cannot represent, with regression coverage for these shapes? These are source-derived cases, not executed queries.

[Visorgood]

Thanks @sunchao – both confirmed against Spark's source. doGenCode evaluates the replacement first while eval checks the array first, and ConstantFolding leaves a throwing expression unfolded with foldable still true, so my !_.foldable test skipped the guard exactly where it was needed.

I've changed the approach rather than patching it. Reproducing Spark's short-circuiting natively is open-ended – two rounds each found a shape I had missed, and since eval and doGenCode disagree, no single guard ordering matches both.

So I took your suggestion of using the dispatcher for what cannot be represented and made it the rule. getSupportLevel now reports Compatible only when the delimiter and null replacement are each a literal or a column read: those cannot throw or carry state, so evaluation order is unobservable and the problem disappears instead of needing guards. Everything else runs through the dispatcher, which is doGenCode itself. convert keeps one guard, IsNull on a nullable replacement, which is the actual #3178 fix. guardedArgs is gone and foldable no longer enters the decision. All four of your cases resolve without a guard.

More conservative than the last revision, but still strictly more native than main, where array_join never runs natively at all.

[sunchao]

Re-reviewed f1ae91a4 against ef62b463, including the updated base. The seven-file PR patch is unchanged, and I found no remaining P1/P2 introduced by this PR. The earlier ordering and stateful-argument concerns are addressed on the default path by limiting native delimiter/replacement arguments to literals or column reads and dispatching the complete expression otherwise.

The incoming dispatcher change does not alter the reviewed Spark 3.5/4.0 ArrayJoin null-short-circuit decision. The throwing foldable-delimiter case returns to the baseline dispatcher route. This does not establish that its reported pre-existing compiler failure is fixed.

This is a source review. I did not run Spark/Comet tests or benchmarks, and the updated multi-version pass counts are author reports. CI, Delta Contrib Build Gate, and CodeQL are still awaiting workflow approval, with no check-run results for this head.

[Visorgood]

Thanks @sunchao! @andygrove, the three test items from your review are addressed; the only one outstanding is CI, which still has no check runs across any push – all four workflows are at action_required on f1ae91a4. Could you approve the workflows so the 3.4 and 3.5 jobs you asked for can actually run?

[sunchao]

The design is sound, but I found one measurable performance improvement worth making: reverse the nullable-replacement guard. I reviewed f1ae91a4 against ef62b463 with five independent scopes.

The current guard construction produces:

IF(replacement IS NULL, NULL, array_to_string(...))

For batches containing both null and non-null replacements, DataFusion filters the referenced input columns into both branches. That copies array/string data for the branch that immediately returns NULL.

Use this equivalent form:

IF(replacement IS NOT NULL, array_to_string(...), NULL)

DataFusion recognizes ELSE NULL, filters only the rows requiring the join, and scatters their results back into place. This preserves short-circuiting and requires only changing the predicate and swapping branches.

I measured both forms using the exact Comet IfExpr, with ten randomized paired runs per case:

Input shapeNull replacementsCurrentReversed
4 elements, short strings	50%	1.156 ms	0.819 ms
4 elements, short strings	90%	0.572 ms	0.186 ms
32 elements, approximately 64-byte strings	90%	25.994 ms	0.659 ms

These are median native expression times per 8,192-row batch, not whole-query results. The improvement reproduced in a second sample. Entirely non-null batches showed no meaningful difference. Results, methodology.

The other design choices look appropriate:

• Keep the small literal/column allowlist. deterministic and foldable do not prove that an expression cannot throw. A broader safety-analysis framework would add substantial complexity. Safe computed delimiters remaining on the dispatcher represent additional optimization opportunities; the base already dispatched those cases.
• Keep the correction in Scala serialization. Reusing the existing conditional and native function is a suitable abstraction. A custom Rust expression or protobuf extension is unnecessary here.
• Do not replace the guard with a post-result null mask. Computed array arguments can throw or carry state. Evaluating them before masking would change which rows execute them.
• Do not add caching for the duplicated replacement expression. Serialization occurs during planning, and the default native path only duplicates cheap literal or column reads. The common two-argument and non-null literal-replacement calls already avoid the guard.

The remaining performance evidence I would request is a focused benchmark comparing head’s default native execution against base’s default JVM dispatcher, using short/long arrays and null-free, mixed-null, and all-null replacements. The component measurements establish the guard improvement, but do not establish that enabling native execution improves whole-query performance.

Validation: 80 batches matched both guard forms and a row-by-row oracle. An independent probe verified identical state and error behavior across successive batches. Current CI has 65 successful and 9 skipped checks. I did not run the JVM suites or an end-to-end Spark benchmark locally.

[sunchao]

Merged, thanks @Visorgood for the contribution! Could you take a look at the above comment and see whether it is worth to add a small follow-up?

[Visorgood]

Thanks @sunchao!

I confirmed the mechanism: CaseExpr::try_new normalizes a null-literal else_expr to None, which takes the single-filter-plus-scatter path instead of materializing both branches. The swap is semantically exact, so this looks well worth doing.

Since #3178 is closed and this is a performance change, I suppose it needs a new issue for the changelog, but just to confirm, shall I file one and open a PR for that new issue?

Agreed the per-batch numbers do not establish a whole-query result, and TPC-H does not exercise array_join, so that would need a targeted benchmark. Happy to look at it separately.