perf: vectorize the native map lookup behind element_at and GetMapValue - #5806
Conversation
`GetMapValue` and `element_at(<map>, key)` both serialise to the DataFusion `map_extract` UDF, and the planner then unwrapped its one-element list with a second `ListExtract` pass. `general_map_extract_inner` re-slices the query key and every candidate key into a fresh `ArrayRef` per comparison and compares them through `dyn Array` equality, which made the lookup roughly 35x more expensive than any other Comet map kernel and slower than Spark itself. Add `SparkMapExtract`, registered under the same `map_extract` name so it overrides the DataFusion one. It runs a single Arrow `eq` over the batch's map entries, scans the resulting bitmask for each row's first match, and gathers the values with one `take`. It also returns the value directly rather than a one-element list, so the `ListExtract` wrapper in the planner goes away. Semantics are unchanged: first matching entry wins, and a missing key, a NULL map row and a NULL lookup key all yield NULL. Key types whose Spark equality the native lookup cannot reproduce are still declined by `MapKeySupport`. `eq` rejects nested key types, so an element-wise comparison remains as a backstop for those. Fixes apache#5795
github-actions
Bot
added
enhancement
|
Thanks @andygrove do you think this recent PR can also help apache/datafusion#24999? |
|
Reviewed against the PR head Overall this is a well-built change. The approach is right, the Arrow usage is correct, and the two hazards specific to this kernel (sliced Major1. The map's validity buffer is never consulted, so a NULL map row is answered from its entriesWhere: The per-row scan derives NULL purely from "no matching entry": let found = (start..end).find(|&i| matched.value(i));
if let Some(i) = found { indices[row] = (i + entries_start) as u32; }
nulls.append(found.is_some());
Spark's This is not a regression against the pinned DataFusion 55.0.0, which has the same gap ( I traced the producers that could deliver such an input (arrow-rs Parquet list reader, arrow-java Suggested direction: one line at the end of the scan, before building let nulls = NullBuffer::union(map_array.nulls(), nulls.finish().as_ref());
Minor2. A constant map argument is expanded to one copy per row before the compareWhere:
This is reachable: If you take it, the cheap version is a scalar-map branch that compares the key array against the single map's 3.
|
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Correctness
Prior state and Spark compatibility
Reviewed head 4f6a1cee80e1491e5d38565e1e32789382967cae against base 5627ab8c01473b186db2fbe73c91c599b1387f9b. Previously, map lookup used DataFusion 55's per-entry array-slice comparisons, returned a one-element list, and required a second ListExtract expression. This PR replaces that path with a value-returning Comet UDF.
The maintained Spark branch-3.5-openai and branch-4.0-openai implementations use ordered key equality, select the first matching entry, and return NULL for missing keys, null lookup keys, null maps, and null matched values. Missing map keys return NULL in both ANSI modes. The new mask scan preserves first-match ordering. Nullable gather indices preserve misses, and take preserves value nulls. The visible entry window and conversion back to absolute entry indices correctly handle sliced maps. Lookup-key coercion and the existing floating-point, collation, and complex-key fallback restrictions remain in place. The array overload is unchanged.
The existing [P2] null-map validity concern in the existing review comment is valid and remains outstanding. A legal map with offsets [0, 1], validity [false], and physical entry a -> 7 yields 7 for lookup a, because the gather validity only tests whether a match exists. Arrow does not require null rows to have empty entry ranges. I independently checked that pinned DataFusion 55 has the same gap, so this is pre-existing, and I am not adding a duplicate inline. Comet's struct-field helper also preserves child buffers while adding parent nulls. The UDF execution layer supplies no automatic validity mask. This is a source-level counterexample, not a demonstrated end-to-end SQL failure. The existing ANSI nullable-input guard can protect element_at. Direct GetMapValue has no equivalent guard. Combining map validity with the result indices and testing a null row with retained entries would close the concern.
Validation and limits
The native CI job passed all 13 new Rust cases. Expression jobs for Spark 3.4, 3.5, 4.0, and 4.1 succeeded. Logs for those jobs and Spark 4.2 show both new Scala cases passed. These six logs checked out merge 36e7522a, whose parents are the reviewed base/head and whose source tree equals the head tree. The Spark 4.2 job passed 1,345 expression tests, then failed during packaging because downloading ASM's POM reported Network is unreachable. Subsequent integration work did not run. Other CI work was still unfinished at the review refresh. No local native/JVM execution or independent benchmark was performed. Maintained Spark 3.4 and 4.1 source branches were unavailable. Apache CI coverage does not establish compatibility with those maintained branches. No additional P1/P2 finding was verified.
Performance
The change removes the repeated candidate ArrayRef slices and the intermediate list/unwrapping pass. A scalar lookup key uses one broadcast comparison over the visible entries. Per-row keys require an entry-index vector and an expanded key array, followed by the same mask scan and value gather. That expansion costs memory proportional to visible entries, including repeated string bytes. The optimization therefore has a clear allocation tradeoff rather than constant scratch space.
The submitted benchmark compares both kernels on 8,192-row batches with 2, 8, or 32 entries per non-null map and constant or varying keys. The author reports roughly 11–56x kernel speedups and an end-to-end lookup reduction from 278 ms to 211 ms, with a documented drifting Spark baseline. Those are author measurements, not results independently reproduced here. They support the targeted workload, not a universal speedup for wide maps or early-hit distributions. Scalar-map expansion is inherited from the old path and is already discussed in the existing review. I found no additional verified performance regression requiring a separate comment.
Design
Approach and key decisions
Returning the actual value type is a useful simplification: the planner can consume the UDF result directly instead of reconstructing a list-extraction expression. Registering the replacement after DataFusion's nested functions is consistent with both Spark serializers emitting map_extract. The separate scalar-key and per-row-key paths expose the main optimization without changing the serialized expression contract, introducing configuration, or broadening Spark's supported key types.
The implementation keeps representation-sensitive logic in one kernel: determine the visible entry window, compare keys, choose each row's first match, and gather values. The existing null-validity concern belongs at that kernel boundary, where every caller receives the same contract. Otherwise, the planner cleanup, registration, tests, and benchmark form a coherent scope.
Abstraction & complexity
The UDF wrapper earns its role by defining coercion and the new return type. The matching helpers separate Arrow's vectorized comparison from the small nested-type fallback. That fallback is outside the current Spark key support gate, but remains localized and tested. The implementation does not require a new shared map framework or another planner abstraction. The main improvement still needed is the already-discussed validity handling and its focused regression case. No additional abstraction or complexity finding warrants a duplicate comment.
The gather validity only recorded whether a match was found, so a NULL map row still returned the value of a matching entry inside its offset range. Arrow does not require a null row's range to be empty: adding a parent null mask over intact child buffers, which Comet's own struct-field helper does, leaves the entries in place, and the UDF layer supplies no validity mask. Spark returns NULL for a NULL map under both ANSI modes, and only element_at has a nullable-input guard upstream of this kernel, so GetMapValue was exposed directly. The row's map validity is now consulted before the mask scan, which also skips scanning a null row's entries. Two tests cover it, one per key path; both return Some(7) for the null row without the fix.
|
The null-map validity concern is fixed in 7d527c1, and your counterexample was exact. The gather validity only recorded whether a match existed, so it never consulted the map's own null buffer. The kernel now checks it before the mask scan and skips a null row's entries entirely, which is also marginally cheaper than scanning them and discarding the result. I built your shape and confirmed it reads the wrong value before the fix rather than trusting the reading. Both pass with the fix, and all 24 I covered both key paths deliberately, because they fail for slightly different reasons. The constant-key path compares the null row's entry against the broadcast key. The per-row path gathers a key per entry, so the null row's entry is compared against that row's own lookup key, and the mask is set there too. Masking at the gather-index step covers both, which is why I put it there rather than in On why there is no Scala test: this shape is not reachable from Spark SQL. Spark's own map builders give a null row an empty offset range, so an end-to-end I left the wider |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Correctness
The null-map validity P2 in the previous review is addressed by 7d527c148. The kernel now appends an invalid gather index and skips the row's match scan when the map is NULL. This covers both constant and per-row lookup keys, including the reported offsets [0, 1, 3] and validity [false, true]. The retained a -> 7 entry can no longer escape as the result of the NULL row.
I rechecked this against the maintained Spark 3.5/4.0 null short-circuit and map-lookup implementations. NULL maps and lookup keys return NULL, missing keys return NULL in both ANSI modes, and the first matching duplicate remains selected. Non-null rows still use the same entry window, match order and absolute gather index. Arrow slices the map validity alongside the visible rows, so the new row-relative validity check also fits sliced inputs. Existing key coercion, unsupported-key fallbacks and the element_at ANSI input guard are unchanged.
The current native CI job passes both retained-entry regressions and all 15 map-lookup tests. Its overall result is 1,273 passed and five skipped. The job checked out merge 4bf5a628, with parents 424c31aa and the reviewed head. That merge contains one newer base commit than the assigned base 5627ab8c. Its six FIRST/LAST-related paths are disjoint from the eight authored paths, all of which match the reviewed head. I did not treat the whole merge tree as identical to the head.
At the 2026-09-09T21:16:07.169750+00:00 refresh, 50 checks had succeeded, six were skipped, ten were running and one was queued. No failed checks were reported. The native test evidence above was inspected in full, but I have not added fresh Spark consumer-log qualification for this update. The Delta build gate passed. I did not run local native/JVM tests or independently reproduce the author's pre-fix failures. Maintained Spark 3.4/4.1 source remains unavailable, so no compatibility claim is made for those branches. No new or remaining verified P1/P2 finding was identified.
Performance
The update adds a map-validity lookup per row and skips the per-row mask scan for NULL rows. It reuses the existing gather-index and validity buffers. It does not add another entry-sized allocation or change the scalar-key comparison and per-row key-expansion paths. The batch comparison still includes retained entries underneath null rows, so this should not be described as avoiding all work for those entries.
The benchmark source is unchanged. Its constant/per-row key cases and 2/8/32-entry maps remain useful for the original optimization, but the reported timings predate this guard and do not measure its incremental cost. No fresh benchmark or additional speedup is claimed here. The existing scalar-map expansion tradeoff is unchanged, and I found no new material performance concern from this correction.
Design
The fix belongs at the shared kernel boundary, where both map lookup callers receive the same NULL contract. Checking validity before scanning each row is sufficient and keeps the correction independent of the Scala guard. The two new tests deliberately construct a legal Arrow representation that ordinary map builders usually normalize away. They directly exercise the reported failure without requiring a new planner path.
Abstraction & complexity
The update introduces no new abstraction or public API. The small fixture isolates the retained-entry representation, and both tests use the existing extraction helper. The correction stays within the existing comparison-and-gather design. I found no additional complexity concern in the update.
|
Lint is stuck would be fixed in #5782 |
comphead
left a comment
comphead
left a comment
There was a problem hiding this comment.
Reviewed for Spark semantics, the JVM/native boundary and the vectorization. No blockers. I traced the kernel against sliced maps, NULL map rows, NULL and per-row lookup keys, duplicate keys, NULL values, empty batches and scalar maps, and found no correctness defect.
Verified along the way:
map_extractis serialized from exactly two places (CometMapExtract,CometElementAt's map branch), so changing the return shape is safe, andregister_all_comet_functionsruns afterfunctions_nested::register_all, so the override takes effect.scalarFunctionExprToProtosets noreturn_type, socreate_scalar_function_exprtakes thecoerce_typesbranch and the result field stays nullable. DroppingListExtractdoes not change the type the JVM sees.- Spark 3.5 / 4.0 / 4.1
GetMapValueUtil.getValueEvalall return NULL on a miss, on a NULL value, and on either NULL operand, with nofailOnErroronGetMapValuesince SPARK-40066. The kernel matches on every axis, and the float / collation / complex-key divergences stay fenced off by the unchangedMapKeySupport. - The diagnosis checks out against the vendored dependency:
datafusion-functions-nested-55.0.0does contain the per-comparisonkeys.slice(start + i, 1)loop.
Side effect worth noting: get_map_value.sql runs its m['a'] cases as query spark_answer_only, so nothing previously asserted that GetMapValue executes natively. The new checkSparkAnswerAndOperator case does.
One correction to the description: "a lookup key whose runtime type is not the map's key type is now rejected rather than silently missing every row". DF 55.0.0's map_extract_inner already rejects it ("The key type {} does not match the map key type {}"). The new check is a different message and place, not new behaviour.
Comments inline. Nothing needs to block a merge. I would take the first two first, both are a few lines.
| let offsets = map_array.offsets(); | ||
| let entries_start = offsets[0] as usize; | ||
| let entries_end = offsets[num_rows] as usize; | ||
| if entries_start == entries_end { |
There was a problem hiding this comment.
Argument validation is now data-dependent: this early return runs before the key-length check (L158) and the key-type check (L223). DF 55.0.0's map_extract_inner validates the key type before touching data, so a mismatch failed identically on every batch. Now a batch whose maps are all empty or all NULL returns NULLs, while a later batch with entries errors. Same query, failure depends on which partition holds data.
Both checks are cheap and neither needs the entries window, so they can be hoisted above this return.
There was a problem hiding this comment.
Fixed. Both checks moved into a validate_lookup_key helper that runs before the empty-window return, so the lookup key type and, for an array key, its length are now checked on every batch regardless of what the maps hold. argument_checks_do_not_depend_on_the_data pins it with an all-empty/all-NULL map: it now errors on a mismatched key type and on a short key array where it previously answered NULLs.
| // `eq` rejects nested key types. `MapKeySupport` declines those before they reach the | ||
| // native lookup, but keep DataFusion's element-wise comparison as a backstop so this | ||
| // kernel is never less capable than the one it replaces. | ||
| Err(_) => Ok(elementwise_match_mask(keys, lookup, lookup_is_scalar)), |
There was a problem hiding this comment.
This catches any ArrowError, not just the nested-type rejection the comment describes. A future or unrelated eq failure silently drops into elementwise_match_mask, which is exactly the per-row slice-and-compare pattern this PR exists to remove. That is a silent ~50x throughput cliff with no signal, and an untestable branch.
Deciding from the type up front would keep genuine errors visible:
if keys.data_type().is_nested() { elementwise_match_mask(..) } else { eq(..)? }There was a problem hiding this comment.
Taken. It dispatches on keys.data_type().is_nested() now, which is exactly the predicate compare_op applies after unwrapping one dictionary level, and with the length and type checked up front nesting was the only remaining thing eq rejects. Both eq calls use ?, so a genuine failure surfaces instead of silently taking the per-row path.
|
|
||
| mod map_extract; | ||
| mod map_sort; | ||
| pub use map_extract::{spark_map_extract, SparkMapExtract}; |
There was a problem hiding this comment.
spark_map_extract has no caller outside this module. comet_scalar_funcs.rs and the new bench both import only SparkMapExtract. (spark_map_sort is exported because comet_scalar_funcs.rs calls it directly.)
The tests are in the same file, so the free function can stay private and only the UDF needs exporting.
There was a problem hiding this comment.
Done. spark_map_extract is private again; only SparkMapExtract is re-exported from map_funcs and lib.rs.
| Arc::new(ScalarUDF::new_from_impl(SparkMakeTime::default())), | ||
| // Overrides datafusion-functions-nested' `map_extract` with a vectorized lookup that | ||
| // returns the value itself rather than a one-element list (#5795). | ||
| Arc::new(ScalarUDF::new_from_impl(SparkMapExtract::default())), |
There was a problem hiding this comment.
DF's MapExtract declares aliases: ["element_at"], and SessionState::register_udf inserts one entry per alias. After this override, udf("map_extract") returns the Comet kernel but udf("element_at") still returns DF's list-returning one.
Nothing serializes that name today, so it is latent rather than live, but the registry is now inconsistent and the next element_at serde would silently get the wrong shape. Either add aliases() returning ["element_at"], or note here why the alias is deliberately left alone.
There was a problem hiding this comment.
Added the alias rather than documenting the gap, since it is three lines and leaves nothing to rediscover. SparkMapExtract::aliases() returns ["element_at"], so the override now replaces both registry entries. I re-grepped and nothing emits that name today, so this is consistency rather than a fix, and there is a one-line test so it does not quietly come undone.
| // a null row can carry a live `a -> 7` that would otherwise match. Spark returns NULL for a | ||
| // NULL map under both ANSI modes, for `element_at` and for `GetMapValue` alike, and only | ||
| // `element_at` has a nullable-input guard upstream of this kernel. | ||
| let map_nulls = map_array.nulls(); |
There was a problem hiding this comment.
Question about reachability, not about the fix.
child_with_parent_nulls only adds a mask, it does not create entries, so this needs a producer that hands Comet a null struct row whose map child already spans a non-empty range. Everything I checked leaves it empty: the Parquet readers emit length 0 for a null parent, Spark's ArrowWriter.StructWriter.setNull recurses into children, and arrow's filter / take / concat compact.
Keep the guard regardless. DF 55.0.0 has the same hole (general_map_extract_inner never consults map_array.is_valid and builds its ListArray with None nulls), so this closes a pre-existing latent bug cheaply, and both tests do fail without it.
The ask is narrower: if you know the producing path, an end-to-end case (SELECT s.m['a'] over a nullable struct<m: map<..>>) would be much stronger than a hand-built array. If you do not, consider softening "Comet's struct-field helper ... leaves the entries in place", which reads as a live path.
There was a problem hiding this comment.
I could not find a producing path either, so I softened the wording rather than claim one. The comment now says every producer traced -- the Parquet readers, Sparks ArrowWriter, and arrows own filter / take / concat -- gives a null row an empty range, and that this is a representation the format permits rather than one known to arrive here, so the guard exists to pin the contract at the kernel boundary. The test fixture`s doc comment had the same phrasing and got the same treatment.
| // MapArray's original entry offsets, so the visible entries start part way into the keys child -- | ||
| // the same trap `mapsort` hit below. Reading the mask from index 0 would answer every row with | ||
| // some other row's entries. | ||
| test("element_at on a sliced map reads the visible entries") { |
There was a problem hiding this comment.
This does exercise entries_start != 0 today, but incidentally.
_2 is a nullable Parquet map column, so under the Spark 4.1 ANSI default needsNullGuard wraps the lookup in CASE WHEN _2 IS NOT NULL, and DF's CaseExpr evaluates the THEN branch through filter_record_batch. It preserves the slice only because the predicate is all-true, so arrow picks IterationStrategy::All and returns values.slice(0, count). Add one NULL map row and the filter compacts the entries, entries_start becomes 0, and the test quietly stops testing what its comment says.
An ANSI-off variant, or _2['a3'] (GetMapValue has no guard) so the sliced map always reaches the kernel, would pin the intent. This is analysis, not an observed failure.
There was a problem hiding this comment.
You were right that it was slicing by accident. It looks up through _2['a3'] / _2['shared'] now, and the fixture carries NULL map rows so the distinction is load-bearing rather than latent; I kept one element_at column for that serde path. I checked that it still detects the bug rather than assuming: replacing keys.slice(entries_start, window_len) with slice(0, window_len) fails the test on the answer, and it passes again with the arithmetic restored.
| /// empty offset range, which is the shape Arrow's builders produce, but nothing in the format | ||
| /// requires it: adding a parent null mask over intact child buffers leaves the entries in | ||
| /// place. Such a row must still read NULL, not the value its live entry holds. | ||
| fn null_row_with_retained_entries() -> MapArray { |
There was a problem hiding this comment.
Minor: the key_field / value_field / StructArray / entries_field / try_new scaffolding is repeated four times (here, map_from, non_string_keys, nested_key_falls_back_to_elementwise_comparison). One map_of(keys, values, offsets, nulls) -> MapArray helper would cut roughly 60 lines and make the three variants read as the variations they are.
There was a problem hiding this comment.
Done. One map_of(keys, values, offsets, nulls) builds all four now, so each fixture is just its own key/value types, offsets and null mask.
| //! Benchmarks for the map lookup behind `GetMapValue` and `element_at(<map>, key)`. | ||
| //! | ||
| //! Each shape is run against both Comet's `SparkMapExtract` and the | ||
| //! `datafusion-functions-nested` `map_extract` it overrides, so the gap that motivated |
There was a problem hiding this comment.
Worth a line here: DF main has already rewritten general_map_extract_inner to a single make_comparator over the batch, so the per-comparison slicing is 55.0.0-specific. This gap will narrow a lot at the next DF bump and the comparison will start measuring something different.
That does not change the case for the Comet kernel (one eq plus one take, plus the removed ListExtract pass), but readers should not take these ratios as permanent.
There was a problem hiding this comment.
Added. The module doc now says to read the ratio as a measurement of the pinned 55.0.0 rather than a permanent gap, names the make_comparator rewrite on DataFusion main as the reason the baseline arm will get much faster at the next bump, and notes that what survives it is the rest of the case for the kernel: one eq plus one take, and the removed ListExtract pass.
Argument validation no longer depends on the data. The lookup key's type and length are properties of the call, so they now run before the empty-entries fast path; previously a batch whose maps were all empty or all NULL answered NULLs while the next batch of the same query errored. The nested-key backstop dispatches on the key type rather than on an `eq` failure. With the length and type already checked, nesting is the only thing `eq` still rejects, so a catch-all `Err(_)` could only hide a genuine error behind a silent per-row throughput cliff. `SparkMapExtract` now declares the `element_at` alias DataFusion's `map_extract` declares. `register_udf` inserts one registry entry per alias, so without it the override was partial and `element_at` still resolved to the list-returning kernel. Nothing in Comet emits that name, so this is consistency, not a fix. `spark_map_extract` is no longer exported: only the UDF has a caller outside the module. The sliced-map Scala test now looks up through `_2[k]`. `element_at` on a nullable operand is wrapped in `CASE WHEN _2 IS NOT NULL` under ANSI, and DataFusion's CaseExpr runs the THEN branch through `filter_record_batch`, which compacts the entries child and resets the first offset to 0 as soon as any row is NULL -- so the old test only sliced by accident. `GetMapValue` has no such guard, and the fixture now carries NULL rows to keep that honest. Reverting the `entries_start` arithmetic still fails the test. `element_at_map.sql` covers the admitted key types the fixtures missed: boolean, tinyint, smallint, bigint, decimal, date, timestamp and timestamp_ntz. Arrow's `eq` is stricter about the exact Arrow type than the `ArrayData` equality it replaced, so a disagreement between it and `coerce_types` would surface as a query failure. A narrower decimal key is not expressible: Spark rejects it at analysis with MAP_FUNCTION_DIFF_TYPES. Also: one `map_of` helper for the four test fixtures, a note that the benchmark ratio measures pinned DataFusion 55.0.0 rather than a permanent gap, and comment corrections for the `DataType::Null` capability this drops and for the null-row-with-retained-entries shape, which no traced producer emits.
|
Thanks @comphead — all eight taken in 68735f0, replies inline. Three of the earlier round's items are in there too. On the constant map argument (finding 2 in your first pass) I went with the comment, but the reachability turned out to be stronger than "deliberate and bounded": the scalar branch is not reachable from Comet at all. The native Key type coverage (finding 5): Description corrected on both counts. The key-type rejection claim is gone, and the two things that are new are stated instead: a NULL lookup key no longer matches a NULL stored key (your finding 7, which is the one real semantic win here), and a Local on this revision: 17 |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed 7d527c14 → e54fdb91. The update validates lookup arguments before empty/null shortcuts, dispatches nested-key comparison explicitly, restores the element_at alias, and strengthens the sliced-map and admitted-key-type tests. The earlier P2 null-map issue remains fixed: both retained-entry regressions pass in the current native test job. No new or remaining verified P1/P2 findings.
CI tested merge f0b7bd5f, whose tree exactly matches this head. All 17 native map lookup tests pass. The Spark 3.5 and Spark 4.0 jobs pass the sliced-map, per-row-key and expanded map SQL coverage using the matching native-library artifact. At the fresh 2026-09-11 06:07 UTC check, checks were 70 successful, 8 skipped. I ran no local native/JVM tests or benchmarks. The author's mutation test remains author-reported. Maintained Spark 3.4/4.1 source branches were unavailable, so I make no compatibility claim for those versions. The incremental code adds no entry-sized allocation. Existing per-row-key/scalar-map allocation costs and benchmark-baseline qualifications remain unchanged.
3 participants
Which issue does this PR close?
Closes #5795.
Rationale for this change
GetMapValue(m[k]) andelement_at(<map>, k)both serialise to the DataFusionmap_extractUDF, and the planner then unwrapped its one-element list result with a second
ListExtractpass.The diagnosis in the issue holds.
general_map_extract_innerre-slices the query key and everycandidate key into a fresh
ArrayRefper comparison and compares them throughdyn Arrayequality:
so the lookup is O(rows x entries-per-row) allocations, and the constant key is not hoisted out of
the row loop. That is why the map lookup was the one map operation where Comet lost to Spark.
What changes are included in this PR?
A native
SparkMapExtract, registered under the samemap_extractname so it overrides theDataFusion one, in
native/spark-expr/src/map_funcs/map_extract.rs:eqover the whole batch of map entries (broadcasting the key when it is a constant),then a bitmask scan for each row's first match, then one
taketo gather the values;ListExtractwrapperin
planner.rsis no longer needed and that second pass goes away;MapArraykeeps its original entry offsets, so only the visible entry window iscompared.
Semantics are unchanged: the first matching entry wins, and a missing key, a
NULLmap row and aNULLlookup key all yieldNULL. Key types whose Spark equality a native lookup cannot reproduce(floating point, non-default collations, complex keys) are still declined by
MapKeySupport, sothey never reach the kernel. Arrow's
eqrejects nested key types, so DataFusion's element-wisecomparison is kept as a backstop for those. The kernel also carries the
element_atalias thatDataFusion's
map_extractdeclares, so the override replaces both registry entries rather thanjust the one Comet emits.
Two behaviour differences against DataFusion 55, in opposite directions:
NULLlookup key no longer matches aNULLstored key. DF 55 compares the two throughArrayDataequality, which reports them equal, so a map carrying aNULLkey -- reachablethrough the
map_from_arraysgap in [Bug] map_from_arrays / map_from_entries do not enforce null-key rejection or spark.sql.mapKeyDedupPolicy #4680 -- returned that key's value where Spark returnsNULL. Upstream'smake_comparatorrewrite does not close this either, sinceSortOptions::default()also treats two nulls as equal.DataType::Nullfirst argument is now an error rather than aNULLpassthrough. Spark'sExtractValue.applyrequires aMapTypechild andCometElementAt.getSupportLeveldeclinesanything that is neither an array nor a map, so this is not reachable from Comet; it is noted
because it is the one place the kernel is less capable than the one it replaces.
Dropping the
ListExtractwrapper also decouples Comet from apache/datafusion#24999, whichchanges the absent-key result from a one-element
NULLlist to an empty list.ListExtractwithfail_on_error=falsewould have absorbed that, but not depending on it is better.How are these changes tested?
Existing tests.
CometMapExpressionSuite(25),CometArrayExpressionSuite(64),CometExpressionSuite(141) and the fullCometSqlFileTestSuite(474, includingelement_at_map.sql,element_at_map_ansi.sql,element_at_map_collation.sqlandget_map_value.sql) all pass.New tests. 17 Rust unit tests covering hit/miss, duplicate keys,
NULLmaps (including aNULLrow whose entries were retained rather than dropped),NULLvalues,NULLand per-rowlookup keys, sliced maps, empty input, a scalar map argument, non-string keys, the nested-key
backstop, the rejected argument types, and that the argument checks do not depend on the data.
Two new
CometMapExpressionSuitecases cover the two tricky paths end to end: a map lookup undera native
OFFSET(sliced map) andelement_at/m[k]with a per-row lookup key. The sliced-mapcase looks up through
m[k]rather thanelement_at, becauseelement_aton a nullable operandis wrapped in
CASE WHEN m IS NOT NULLunder ANSI and DataFusion'sCaseExprruns the THEN branchthrough
filter_record_batch, which compacts the entries child and resets the first offset to 0 assoon as any row is
NULL.GetMapValuehas no such guard. Reverting theentries_startarithmetic fails the test.
element_at_map.sqlgains the admitted key types the fixtures were missing -- boolean, tinyint,smallint, bigint, decimal, date, timestamp and timestamp_ntz -- since Arrow's
eqis stricterabout the exact Arrow type (decimal precision and scale, timestamp time zone, integer width) than
the
ArrayDataequality it replaces, so any disagreement withcoerce_typessurfaces as a queryfailure rather than a miss. A narrower decimal lookup key is not expressible: Spark rejects it at
analysis with
MAP_FUNCTION_DIFF_TYPES.Kernel benchmark (new
native/spark-expr/benches/map_extract.rs, 8192-row batches ofmap<string, string>, 10%NULLrows, 60 distinct keys, run against both implementations):datafusioncometEnd to end, reproducing the issue's shape (2,000,000 rows,
attrs map<string, string>, 10%NULLmaps, 0-6 entries each, 60 distinct keys, 400 distinct values;noopsink, median of 5after 2 warmups; Apple silicon, so the absolute numbers differ from the issue's Ryzen run):
SELECT attrs FROM tSELECT size(attrs) FROM tSELECT size(map_keys(attrs)) FROM tSELECT size(map_values(attrs)) FROM tSELECT element_at(attrs, 'a1') FROM tTaking
size(attrs)as the control, as the issue does, the isolated cost of the lookup goes from+64 ms to +1.5 ms, and
element_atis no longer the outlier among the map kernels. TheSpark-arm numbers drifted about 15% between the two runs on this machine, so read the Comet columns
against each other and against their own control rather than the ratios.