fix, perf: Use Arrow comparator for key comparison in map_extract - #24999
Conversation
The previous implementation of `map_extract` did the following for row: 1. Create a one-element array slice containing the row's search key 2. Scan the map's entries. For each entry, create a one-element array slice and compare the two slices using Arrow's array equality 3. Stop at the first match; if no matches, append a NULL instead This had three shortcomings: 1. It was very inefficient, because a lot of allocations are done for every element of every map. 2. It got the equality semantics wrong for some corner-cases. In particular, maps with dictionary-valued keys might encode a logical NULL in two physically distinct ways (apache#24983). Arrow's array equality also considers sparse unions that have different values in unselected child fields to be distinct; this is arguably a bug in Arrow though. 3. It returned `[NULL]` for missing map keys instead of an empty list, which is the behavior implemented by DuckDB (apache#24981). Instead, we can implement `map_extract` with a single arrow-ord comparator. This enables comparing the search key with each map element directly by index, without allocating. It also avoids the differences in comparison semantics outlined above. Finally, this PR fixes the behavior for absent map keys to be consistent with DuckDB. Benchmark results (M4 Max): - int32/first/1024x32, 135.629 µs -> 6.234 µs, -95.40% - int32/last/1024x1, 133.110 µs -> 6.157 µs, -95.37% - int32/last/1024x32, 3477.393 µs -> 38.562 µs, -98.89% - int32/last/1x0, 0.498 µs -> 0.313 µs, -37.23% - int32/last/1x1, 0.593 µs -> 0.384 µs, -35.34% - int32/missing/1024x32, 3472.698 µs -> 34.889 µs, -99.00% - int32/varying/1024x32, 1844.686 µs -> 25.307 µs, -98.63% - struct/first/1024x32, 335.421 µs -> 8.769 µs, -97.39% - struct/last/1024x1, 335.097 µs -> 8.647 µs, -97.42% - struct/last/1024x32, 8325.830 µs -> 71.611 µs, -99.14% - struct/last/1x0, 0.466 µs -> 0.242 µs, -48.08% - struct/last/1x1, 0.747 µs -> 0.391 µs, -47.67% - struct/missing/1024x32, 8451.153 µs -> 61.480 µs, -99.27% - struct/varying/1024x32, 4498.016 µs -> 41.343 µs, -99.08% - utf8_view/first/1024x32, 218.809 µs -> 9.877 µs, -95.49% - utf8_view/last/1024x1, 190.468 µs -> 8.464 µs, -95.56% - utf8_view/last/1024x32, 6016.140 µs -> 124.405 µs, -97.93% - utf8_view/last/1x0, 0.526 µs -> 0.353 µs, -32.97% - utf8_view/last/1x1, 0.762 µs -> 0.523 µs, -31.42% - utf8_view/missing/1024x32, 5999.511 µs -> 114.082 µs, -98.10% - utf8_view/varying/1024x32, 3226.583 µs -> 71.732 µs, -97.78% ("1024x32" means 1024 rows and each row is a map with 32 entries.)
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24999 +/- ##
==========================================
+ Coverage 81.67% 81.72% +0.04%
==========================================
Files 1126 1127 +1
Lines 414842 416390 +1548
Branches 414842 416390 +1548
==========================================
+ Hits 338841 340291 +1450
- Misses 56070 56101 +31
- Partials 19931 19998 +67 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
jayzhan211
left a comment
There was a problem hiding this comment.
Thanks @neilconway,
The kernel looks right to me — I checked the slice handling specifically (MapArray::slice keeps entries unsliced with absolute offsets, and this code correctly indexes keys/values absolutely while indexing query_keys_array and nulls by slice-relative row), and probed a few cases the tests don't cover: a null map row with a non-empty offset range, an all-empty-map array with a null buffer, and a matched key whose value is NULL under use_nulls: false. All correct.
The gap is documentation. This changes map_extract/element_at results in three user-visible ways, and one isn't mentioned anywhere:
- missing key:
[NULL]→[](announced, #24981) - NULL map row:
[NULL]→NULL— not in the PR description or either linked issue (map.slt:691, map.slt:735) - NULL lookup key:
[NULL]→[](map.slt:655)
All three are silent: nothing errors, results just differ. array_length(map_extract(m, k)) goes from always-1 to 0/1/NULL, so downstream filters built on the old shape change meaning with no signal.
Could you add an entry to docs/source/library-user-guide/upgrading/56.0.0.md (it already carries this kind of change — see the GroupColumn and datafusion-proto sections), and mention #2 in the PR description? Something like:
### `map_extract` / `element_at` return an empty list for absent keys
`map_extract` (and its alias `element_at`) previously returned a single-element
list containing NULL when the key was not present in the map. It now returns an
empty list, matching the documented behavior and DuckDB. Related changes:
- A NULL map input row now yields NULL instead of `[NULL]`.
- A NULL lookup key now yields `[]` instead of `[NULL]`.
**Migration guide:**
```sql
-- Before
SELECT map_extract(MAP {'a': 1}, 'missing'); -- [NULL]
SELECT array_length(map_extract(MAP {'a': 1}, 'missing')); -- 1
-- After
SELECT map_extract(MAP {'a': 1}, 'missing'); -- []
SELECT array_length(map_extract(MAP {'a': 1}, 'missing')); -- 0
Expressions that relied on the result always having length 1 should switch to
checking array_length(...) = 0 (or cardinality) for the absent-key case.|
@jayzhan211 Thanks! I added a migration guide entry and updated the PR description. I also added another unit test. |
|
run benchmark map |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing neilc/fix-map-extract-comparator (e6525b5) to 262936e (merge-base) diff Run configurationrun benchmark mapResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing neilc/fix-map-extract-comparator (e6525b5) to 262936e (merge-base) diff Run configurationrun benchmark mapCPU Details (lscpu)Details
Resource Usagemap — base (merge-base)
map — branch
File an issue against this benchmark runner |
comphead
left a comment
There was a problem hiding this comment.
Thanks @neilconway it looks promising, a small nit you may want to address is
Third copy of the map-key scan loop; get_field keeps the slower one
Description: datafusion/functions/src/core/getfield.rs:107-153
(process_map_array) and :159-194 (process_map_with_nested_key) run the same
"scan a map row's entries, take the first key match, emit the value" loop that
general_map_extract_inner now runs. After this PR the three copies have diverged:
map_extract sizes MutableArrayData by output length with use_nulls: false,
while both getfield.rs copies still use
Capacities::Array(original_data.len()) with use_nulls: true, and
process_map_with_nested_key re-reads map_array.value_offsets() inside the loop.
Reason: map[key] is the more common surface than map_extract(map, key), and
it does not get any of this PR's improvement. The next map-key bug will need three
fixes.
|
@comphead Thanks! That is a helpful suggestion. I poked around at the |
|
@jayzhan211 @comphead thanks for the reviews! |
Which issue does this PR close?
map_extractfails to match equal struct keys with differently encoded dictionary nulls #24983map_extractreturns[NULL]instead of[]for missing keys #24981Rationale for this change
The previous implementation of
map_extractdid the following for each row:This had three shortcomings:
map_extractfails to match equal struct keys with differently encoded dictionary nulls #24983).[NULL]for missing map keys instead of an empty list (map_extractreturns[NULL]instead of[]for missing keys #24981); returning an empty list is what the DataFusion docs claim this function does, and it is the DuckDB behavior.Instead, we can implement
map_extractwith a single arrow-ord comparator. This enables comparing the search key with each map element directly by index, without allocating. It also avoids the differences in comparison semantics outlined above.Finally, this PR fixes the behavior for absent map keys to be consistent with DuckDB.
Benchmark results (M4 Max):
("1024x32" means 1024 rows and each row is a map with 32 entries.)
What changes are included in this PR?
map_extractto use a comparatormap_extractmap_extractfails to match equal struct keys with differently encoded dictionary nulls #24983)map_extractreturns[NULL]instead of[]for missing keys #24981)What is the testing strategy for this PR?
Existing tests pass; new tests added. Verified that the new tests fail if the implementation is reverted.
Are there any user-facing changes?
Yes, semantics of
map_extracthave changed in the following cases:[]for an absent key, rather than[NULL]NULLwhen called on a map that isNULL, rather than[NULL][]when called with aNULLkey, rather than[NULL]In all three cases, the new behavior matches DuckDB.