fix: dispatch map lookups with normalized keys and nondeterministic null-guarded children

[dwsmith1983]

Which issue does this PR close?

Closes #5580, closes #5781.

Rationale for this change

Two serde gaps in the array and map expressions, both resolved by the codegen dispatcher rather than by native changes.

map_col[key] and element_at(map, key) decline float, collated and complex map keys because the native lookup compares raw Arrow values where Spark finds a 0.0 key with -0.0, treats NaN as equal to itself, compares strings by collation and compares complex keys with interpreted ordering. The declines are right, but neither serde mixed in CodegenDispatchFallback, so the whole projection fell back to Spark instead of running Spark's own generated code inside the Comet pipeline.

size, array_append, arrays_zip and map_from_arrays reproduce Spark's NULL propagation with a CASE WHEN child IS NOT NULL guard that serializes the child twice. A stateful child advances each copy independently, so the guard and the operation see different rows and the answer is silently wrong: on a 16-row table with IF(monotonically_increasing_id() % 2 = 0, array(1), NULL) as the operand, size returned -1 on five rows where Spark returns 1, arrays_zip returned [null, 2] for [1, 2], array_append returned [2] for [1, 2] and map_from_arrays returned NULL for {1 -> 2}. element_at had the same shape and was fixed in #5766 for its ANSI arm; these four are not ANSI-gated, so the wrong answers were reachable in every configuration.

What changes are included in this PR?

• CometMapExtract and CometElementAt mix in CodegenDispatchFallback, so the declined key types run through the dispatcher. CometElementAt's ANSI arm for a nondeterministic operand dispatches the same way.
• A shared NullGuardSupport gate declines any nondeterministic child in CometSize, CometArrayAppend (the array and the item, which sits inside the guard's THEN branch), CometArraysZip and CometMapFromArrays, and all four mix in CodegenDispatchFallback, so Spark's generated code evaluates the child once. Nullability is not consulted: a non-nullable stateful child only stays correct today because DataFusion skips the filter when the guard matches every row, which is not a contract to rely on.
• CometArrayAppend reports a nullable array with an item other than a literal or column as Incompatible, whatever the ANSI setting, so that shape runs through the dispatcher by default. The native guard skips the item on rows whose array is NULL, and an item such as element_at(array(1), _1 - 1) raises there in Spark with ANSI off too. A literal or column item cannot raise and stays native. The literal-or-column check is the one CometArrayContains and CometArrayJoin already used, now shared as EvaluationOrder.isUnobservable.
• With subexpression elimination on, Spark's projection evaluates a subexpression that several of its expressions share for every row, before their NULL checks, while the dispatcher compiles each expression alone and can skip it on a row where an earlier argument is NULL. QueryPlanSerde.skipsSharedSubexpression finds a node in the tree that can skip an argument holding such a subexpression: any argument after the first, the dividend of a division or pmod (the divisor is evaluated first), the first two arguments of locate (start is evaluated first), and none for Spark 3.x array_append, which evaluates both. The six serdes above decline dispatch in that case and fall back to Spark, as before this change. A subexpression the expression repeats on its own does not count while spark.sql.subexpressionElimination.enabled is on, since the kernel then eliminates it as well. With the setting off, only an aggregate counts its shared subexpressions, like canShortCircuitNulls, because Spark's whole-stage codegen aggregates share the subexpressions of their update expressions even then. A projection shares nothing, so its calls stay dispatched.
• The native shuffle key checks in CometShuffleExchangeExec and the projection check in CometTakeOrderedAndProjectExec.getSupportLevel now run in the same withSharedSubexpressions scope that CometNativeShuffleWriter and finalNativePlan use when they serialize. Before, a decline that only fires inside that scope would let planning choose the native path and then fail at execution. The columnar shuffle checks are unchanged, since that path evaluates its keys with Spark's own projection.
• getUnsupportedReasons lists updated for the doc generator, and the six affected rows in the expressions guide move from Native to Hybrid, and the array_append row says which items are dispatched.

How are these changes tested?

SQL-file fixtures over parquet tables, all asserting Spark's answer and the execution path:

• Four new *_nondeterministic_child.sql fixtures, one per serde, with the stateful operand (dispatched), a non-nullable stateful operand (dispatched), and a deterministic nullable operand (native). array_append also pins that a stateful item is dispatched, and its fixture is capped at Spark 3.5 because 4.x rewrites array_append to array_insert. Before the change, the stateful cases failed as result mismatches with the values above.

• element_at_map.sql and get_map_value.sql flip their fallback cases to dispatch and add NaN lookups and a struct-keyed map column with a per-row key, a NULL inside the key and a NULL key. element_at_map_collation.sql flips the same way and passes on the Spark 4.0 profile. element_at_ansi.sql pins the dispatched nondeterministic arm.

• codegen_dispatch_calendar_interval_output.sql pins that map_from_arrays declines the dispatcher for a calendar-interval result on its float-key and LAST_WIN routes.

• map_builders_float_key_gate.sql, map_builders_float_key_gate_legacy.sql and map_builders_float_key_last_win.sql give the FLOAT call its own values array (array(3, 4)). The SQL test harness turns off ConstantFolding, so an array(1, 2) shared by both calls counts as a shared subexpression and declines the dispatch; in a real query it folds to a literal.

• CometMapExpressionSuite renames five fallback tests to dispatch tests through the helper that checks the dispatch tag.

• array_append_ansi_null_array.sql (ANSI on) and array_append_null_array_item.sql (ANSI off) pin that an item raising on a NULL-array row raises through the dispatcher as in Spark (1 / (_1 - 1) under ANSI, element_at(array(1), _1 - 1) without it), that a column or literal item over a nullable array stays native (the ANSI file adds a NULL literal), that any other item over a nullable array is dispatched, and that a non-nullable array stays native and raises on Comet when the item raises. The ANSI-off file also pairs two calls that share element_at(array(1), _1 - 1) inside map lookups, which fall back to Spark and raise. The arrays_zip fixture adds a dispatched query over an array<int> column so the kernel copies list vectors from real input.

• Six *_shared_subexpression.sql fixtures, with ANSI and subexpression elimination on, pair two expressions that share 1.0D / x over a row with NULL input and x = 0, so Spark raises DIVIDE_BY_ZERO. They cover each of the six serdes, a division and a pmod in the tree, locate with a NULL start, an aggregate, and the operator paths (a hash shuffle, a global sort, and TakeOrderedAndProject). Controls pin that a single call, a column key, a subexpression shared only in the first argument and one the expression repeats itself all stay dispatched. With subexpression elimination off, an aggregate pair still raises like Spark and a projection pair stays dispatched and returns rows like Spark. Each decline, evaluation-order rule and planner scope was removed in turn to confirm a fixture fails without it.

Run on Spark 4.1 (all map and array SQL fixtures, CometMapExpressionSuite, CometNativeShuffleSuite, CometTopKSuite), Spark 3.5 (the shared subexpression and array SQL fixtures, CometArrayExpressionSuite and CometMapExpressionSuite), Spark 4.0 (the shared subexpression fixtures and CometMapExpressionSuite), plus the strict-warnings compile.

Performance

#5875 measured this same routing with CometCodegenDispatchBenchmark. At 65,536 rows, DOUBLE-key lookup-only projections ran 8.7 to 11.4 percent slower through the dispatcher than with it switched off, and mixed projections with a DOUBLE key and a small map ran 10.5 to 12.7 percent faster. The slower lookup-only case is accepted here in exchange for keeping the whole projection in Comet, where before it fell back to Spark.

What this does not cover

array_append keeps its item inside the native NULL guard's THEN branch, which DataFusion evaluates only on the rows where the array is not NULL, while Spark's codegen evaluates the item on every row. A nondeterministic item is declined to the dispatcher for that reason. A deterministic item that raises on a row whose array is NULL, such as element_at with index 0 or a division by zero under ANSI, would return cleanly on the native path where Spark raises. The serde therefore reports a nullable array with an item other than a literal or column as Incompatible, so array_append runs through the codegen dispatcher by default and raises like Spark, and the native path stays available under allowIncompatible with that divergence recorded in the compatibility guide. The check looks at the item's shape, not at whether it can raise, so an arithmetic item, or a column the analyzer wrapped in a widening cast (array_append(bigint_arr, int_col)), is dispatched as well. Evaluating the item outside the guard, so those shapes can run natively again, is tracked in #6086.

The native element_at and map_from_arrays guards skip a shared subexpression the same way on main, tracked in #6970. The other dispatch serdes on main have the gap this PR closes for its six, tracked in #6971. map_from_entries does not yet decline the dispatcher for a calendar-interval result, tracked in #7083.

[andygrove]

The comment above getSupportLevel says only the array operand sits under the NULL guard, but the item is inside the THEN branch too, and DataFusion's CaseExpr filters the batch before it evaluates THEN. So the item is evaluated only on the rows the guard selects.

That matters for array_append specifically, because it is the one expression in this group whose Spark codegen does not short-circuit. On 3.4 and 3.5 ArrayAppend.doGenCode emits leftGen.code + rightGen.code + ctx.nullSafeExec(left.nullable, leftGen.isNull) { ... }, so rightGen.code runs on every row even when the array is NULL. ElementAt and MapFromArrays go through nullSafeCodeGen and Size has one child, which is why the guard matches Spark for those three.

On your own 16 row table, SELECT _1, array_append(IF(_1 % 2 = 0, array(1), CAST(NULL AS ARRAY<INT>)), monotonically_increasing_id()) FROM test_array_append_nondet gives Spark [1,0], [1,2], [1,4] and so on for the even rows, because the counter advances on all 16. Comet evaluates the native monotonically_increasing_id only over the 8 filtered rows and gives [1,0], [1,1], [1,2] up to [1,7]. The same filtering also swallows an error the item would have raised. Under ANSI, array_append(IF(_1 % 2 = 0, array(1), CAST(NULL AS ARRAY<INT>)), 1 / (_1 - 1)) raises DIVIDE_BY_ZERO in Spark at _1 = 1 and returns cleanly in Comet.

The expect_native case below uses array(1), which is non-nullable, so the guard mask is all true and CaseExpr takes the branch that skips the filter entirely. That case passes whatever the serde does, which is why the fixture does not catch this. Would it make sense for NullGuardSupport to cover both children of ArrayAppend, and to change that fixture to a nullable operand with a stateful item so it fails without the change?

On sequencing, I have commented on #5875 pointing it here, since this PR carries the null-guard correctness work as well as the map lookup dispatch and that makes it the better base. Two things to settle before it lands though. #5875 measured this exact routing change and found lookup-only projections running 8.7 to 11.4 percent slower with the dispatcher on at 65536 rows, with no reproducible benefit at 1024 rows. Do you have numbers for the routing here? And #5854 teaches the native map builders LAST_WIN and null key rejection, which would remove the Incompatible branch this PR routes through the dispatcher for map_from_arrays, so that one needs an order agreed with @peterxcli.

Last thing is coverage for the routes this enables. map_from_arrays_dedup_policy.sql sets spark.comet.exec.scalaUDF.codegen.enabled=false and keeps asserting expect_fallback(mapKeyDedupPolicy), so the new LAST_WIN dispatch route is never exercised. Could those queries assert expect_dispatch(map_from_arrays) instead, with a separate dispatcher-off file if the fallback is still worth pinning? The four new nondeterministic fixtures also build every array inline over an int column, so no dispatched kernel in them reads an Arrow ListVector from a column. #5844 hit a real Spark 3.4.3 ColumnarArray.copy problem in exactly that shape, so an array<int> column in at least the size and map_from_arrays fixtures would be worth having. While you are in expressions.md, cardinality still says Native two rows above size, and Spark registers both names against Size.

[dwsmith1983]

Would it make sense for NullGuardSupport to cover both children of ArrayAppend, and to change that fixture to a nullable operand with a stateful item so it fails without the change?

Yes to both. CometArrayAppend now declines when either child is nondeterministic, and the fixture's native case is a nullable array with a stateful item, which fails without the change on exactly the 16-row table you used. The ANSI case where the item raises on a filtered row is the same mechanism but not nondeterminism; I have noted it in the description as the remaining gap for the guard shape rather than widening the decline to every non-literal item here.

Could those queries assert expect_dispatch(map_from_arrays) instead, with a separate dispatcher-off file if the fallback is still worth pinning?

Done: map_from_arrays_dedup_policy.sql asserts the dispatched route with the dispatcher on, and map_from_arrays_dedup_policy_dispatcher_off.sql pins the fallback.

an array<int> column in at least the size and map_from_arrays fixtures would be worth having

Added to both: the dispatched kernel now reads its arrays from columns as well as building them inline. cardinality says Hybrid.

Do you have numbers for the routing here?

Not yet; I will run the same shape #5875 measured before this lands. On #5854, once native LAST_WIN lands the map_from_arrays dispatch branch goes away and only the dispatcher-off fixture needs to follow, so I am fine with #5854 landing first.

[andygrove]

Triage note: #5875 also closes #5580, and mixes CodegenDispatchFallback into the same two serdes (CometMapExtract and CometElementAt in maps.scala) with overlapping element_at / get_map_value fixtures. @LinSimon-901101 flagged the overlap in that description rather than letting us find it in review, and offered to defer.

From reading both, this one looks like the superset: it also closes #5781 for the stateful child being serialized twice under the CASE WHEN child IS NOT NULL guard in size, array_append, arrays_zip and map_from_arrays, which #5875 does not touch. Are you both happy for this PR to carry the map-lookup serde, with any non-overlapping coverage from #5875 rebased on top once it lands?

[dwsmith1983]

Are you both happy for this PR to carry the map-lookup serde, with any non-overlapping coverage from #5875 rebased on top once it lands?

Yes, happy for this PR to carry the map-lookup serde. The size, array_append, arrays_zip and map_from_arrays fix has no counterpart in #5875, so rebasing its extra fixtures and benchmark cases on top once this lands is the smaller move. One sequencing note: if #5854 lands first, its native map_from_arrays replaces the LAST_WIN decline this PR wraps, and the gate here becomes the fallback chain peterxcli described on that PR.

[dwsmith1983]

The [scans] failure on Spark 3.5 is not from this PR. The two DPP - single-split sort-merge join (AQE=true, ...) tests came in with #6814 and fail the same way on main with the 3.5 profile, since Spark 3.5 replaces the pruned join with an empty LocalTableScan rather than an EmptyRelation. Tracked in #6901.

[sunchao]

Reviewed commit d1302c055d26497a81f92d16ebbc2d47cbdea3bb. New P1/P2 findings are detailed inline.

[sunchao]

[P2] [P2] Preserve shared-expression evaluation when dispatching map lookups. With ANSI and spark.sql.subexpressionElimination.enabled=true, a Parquet table containing nullable DOUBLE-key maps m1, m2 and x = 0 makes SELECT element_at(m1, 1.0D/x), element_at(m2, 1.0D/x) FROM t raise DIVIDE_BY_ZERO in Spark even when both maps are NULL. Spark’s ProjectExec evaluates the shared division before the lookup NULL guards. This newly enabled route serializes each lookup separately, so both generated kernels skip the division and return NULL, allowing the query to succeed instead. At the base, DOUBLE-key lookups fall back to Spark. Could this route preserve eager shared-expression evaluation, or retain Spark fallback for this shape? canShortCircuitNulls already identifies the constraint, but dispatch bypasses it.

Evidence: Reproduced with freshly compiled head serializers/planner/dispatcher on Spark 4.1.3 in /tmp/review5867-d1302-current-_xjobhhq/SharedDeterministicProbe.scala, with results in SharedDeterministicProbe.log. Input equivalent: CREATE TABLE t(m1 MAP<DOUBLE,INT>, m2 MAP<DOUBLE,INT>, x DOUBLE) USING parquet; INSERT INTO t VALUES (NULL, NULL, 0.0D);. Spark raises DIVIDE_BY_ZERO for the two-lookup query. Either lookup alone returns NULL, and disabling subexpression elimination makes the pair return [NULL,NULL]. With elimination enabled, the exact-head planner reports canShortCircuitNulls=false yet emits both JVM dispatch expressions. Evaluating those serialized expressions returns NULL for each column. Spark’s ProjectExec.doConsume confirms eager common-expression evaluation. This is a component reproduction, not a full native execution run.

[dwsmith1983]

Could this route preserve eager shared-expression evaluation, or retain Spark fallback for this shape?

It keeps Spark fallback for it. QueryPlanSerde.skipsSharedSubexpression looks, inside the operator's withSharedSubexpressions scope and with subexpression elimination on, for a node in the expression that can skip an argument holding a subexpression the operator shares. For most nodes that is any argument after the first. Division and pmod evaluate the divisor first and skip the dividend, and locate skips its first two arguments when start is NULL. The six serdes this PR newly dispatches decline in that case, so element_at(m1, 1.0D/x), element_at(m2, 1.0D/x) falls back to Spark as it did at the base. A subexpression the expression repeats on its own stays dispatched, since the kernel eliminates it too and evaluates it for every row.

Declining exposed two plan-time checks that ran outside that scope while serialization ran inside it: the native shuffle key check and the TakeOrderedAndProject projection check. Without a fix, DISTRIBUTE BY or a sort on such a pair would have chosen native shuffle and then failed at execution. Both now use the scope the writer and finalNativePlan use.

The native element_at and map_from_arrays guards have the same gap on main with INT keys, which is #6970, and the other dispatch serdes on main have it as well, which is #6971.

[dwsmith1983]

One change since: the check no longer looks at spark.sql.subexpressionElimination.enabled. Spark's whole-stage codegen aggregates share the subexpressions of their update expressions even with it off, so max(element_at(m1, 1.0D/x)), max(element_at(m2, 1.0D/x)) still raises there, and the dispatched lookups would not have. These shapes now fall back whatever the setting says.

[andygrove]

Thanks for the work on this. I went through the current head again, including the calendar-interval declines and the shared-subexpression check, and compared the touched serdes against branch-1.1. The routing changes look right to me, and the fixtures fail without the declines.

One thing is still open, in CometArrayAppend.getSupportLevel (arrays.scala). The SQLConf.get.ansiEnabled && expr.left.nullable gate misses in both directions.

With ANSI off, array_append(IF(_1 % 2 = 0, array(1), CAST(NULL AS ARRAY<INT>)), element_at(array(1), _1 - 1)) should raise INVALID_INDEX_OF_ZERO in Spark 3.4 and 3.5. element_at raises on index 0 whatever the ANSI setting, and ArrayAppend.doGenCode evaluates the item on every row. The native CASE only evaluates the item where the array is not NULL, so Comet returns a value. With ANSI on, a plain column item is dispatched even though it can't raise. I did not rerun this at the current head, so please correct me if I have misread the gate.

Would it make sense to gate on the item instead of on the ANSI setting? CometArrayJoin.orderInsensitive in this file checks for a literal or a column, and that is the case that is safe to keep native. Could array_append_ansi_null_array.sql also gain a case with ANSI off?

[dwsmith1983]

Would it make sense to gate on the item instead of on the ANSI setting?

Yes, you read the gate right, and it is changed in 4fde11b. CometArrayAppend now reports a nullable array with an item other than a literal or column as Incompatible whatever the ANSI setting, so your IF(...) repro runs through the dispatcher and raises INVALID_INDEX_OF_ZERO with ANSI off, and a column or literal item stays native with ANSI on. The literal-or-column check is the one CometArrayContains and CometArrayJoin each had, now shared as EvaluationOrder.isUnobservable.

The ANSI-off cases are in a new array_append_null_array_item.sql, since the ANSI file sets the mode for the whole file. It has your repro and the same item over a nullable column (both expect_error_class), column and literal items native, _1 + 1 dispatched, a non-nullable array staying native and raising when the item raises, and a pair sharing the element_at that falls back. array_append_ansi_null_array.sql now expects the column item native and adds a NULL literal item. Because the check is on the item's shape, a column the analyzer wraps in a widening cast is dispatched too; the description says so.