Skip to content

fix, perf: Use Arrow comparator for key comparison in map_extract - #24999

Merged
neilconway merged 3 commits into
apache:mainfrom
neilconway:neilc/fix-map-extract-comparator
Sep 8, 2026
Merged

fix, perf: Use Arrow comparator for key comparison in map_extract#24999
neilconway merged 3 commits into
apache:mainfrom
neilconway:neilc/fix-map-extract-comparator

Conversation

@neilconway

@neilconway neilconway commented Sep 6, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

The previous implementation of map_extract did the following for each 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 (map_extract fails to match equal struct keys with differently encoded dictionary nulls #24983).
  3. It returned [NULL] for missing map keys instead of an empty list (map_extract returns [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_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.)

What changes are included in this PR?

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_extract have changed in the following cases:

  • We now return [] for an absent key, rather than [NULL]
  • We now return NULL when called on a map that is NULL, rather than [NULL]
  • We now return [] when called with a NULL key, rather than [NULL]

In all three cases, the new behavior matches DuckDB.

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.)
@github-actions github-actions Bot added sqllogictest SQL Logic Tests (.slt) functions Changes to functions implementation labels Sep 6, 2026
@codecov-commenter

codecov-commenter commented Sep 6, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 93.81443% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.72%. Comparing base (262936e) to head (e6525b5).
⚠️ Report is 22 commits behind head on main.

Files with missing lines Patch % Lines
datafusion/functions-nested/src/map_extract.rs 93.81% 0 Missing and 6 partials ⚠️
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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@jayzhan211 jayzhan211 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  1. missing key: [NULL][] (announced, #24981)
  2. NULL map row: [NULL]NULL — not in the PR description or either linked issue (map.slt:691, map.slt:735)
  3. 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.

@github-actions github-actions Bot added the documentation Improvements or additions to documentation label Sep 7, 2026
@neilconway

Copy link
Copy Markdown
Contributor Author

@jayzhan211 Thanks! I added a migration guide entry and updated the PR description. I also added another unit test.

@comphead

comphead commented Sep 7, 2026

Copy link
Copy Markdown
Contributor

run benchmark map

@adriangbot

Copy link
Copy Markdown

🤖 Benchmark running (GKE) | trigger
Instance: c4a-highmem-16 (12 vCPU / 65 GiB) | Linux bench-c5575319922-2208-mvvj4 6.12.94+ #1 SMP Fri Jul 17 09:42:57 UTC 2026 aarch64 GNU/Linux

CPU Details (lscpu)
Architecture:                            aarch64
CPU op-mode(s):                          64-bit
Byte Order:                              Little Endian
CPU(s):                                  16
On-line CPU(s) list:                     0-15
Vendor ID:                               ARM
Model name:                              Neoverse-V2
Model:                                   1
Thread(s) per core:                      1
Core(s) per cluster:                     16
Socket(s):                               -
Cluster(s):                              1
Stepping:                                r0p1
BogoMIPS:                                2000.00
Flags:                                   fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache:                               1 MiB (16 instances)
L1i cache:                               1 MiB (16 instances)
L2 cache:                                32 MiB (16 instances)
L3 cache:                                80 MiB (1 instance)
NUMA node(s):                            1
NUMA node0 CPU(s):                       0-15
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Not affected
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; CSV2, BHB
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Not affected

Comparing neilc/fix-map-extract-comparator (e6525b5) to 262936e (merge-base) diff

Run configuration
run benchmark map

Results will be posted here when complete


File an issue against this benchmark runner

@adriangbot

Copy link
Copy Markdown

🤖 Benchmark completed (GKE) | trigger

Instance: c4a-highmem-16 (12 vCPU / 65 GiB)

Comparing neilc/fix-map-extract-comparator (e6525b5) to 262936e (merge-base) diff

Run configuration
run benchmark map
CPU Details (lscpu)
Architecture:                            aarch64
CPU op-mode(s):                          64-bit
Byte Order:                              Little Endian
CPU(s):                                  16
On-line CPU(s) list:                     0-15
Vendor ID:                               ARM
Model name:                              Neoverse-V2
Model:                                   1
Thread(s) per core:                      1
Core(s) per cluster:                     16
Socket(s):                               -
Cluster(s):                              1
Stepping:                                r0p1
BogoMIPS:                                2000.00
Flags:                                   fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache:                               1 MiB (16 instances)
L1i cache:                               1 MiB (16 instances)
L2 cache:                                32 MiB (16 instances)
L3 cache:                                80 MiB (1 instance)
NUMA node(s):                            1
NUMA node0 CPU(s):                       0-15
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Not affected
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; CSV2, BHB
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Not affected
Details

group                                    HEAD                                   neilc_fix-map-extract-comparator
-----                                    ----                                   --------------------------------
make_map_1000                            1.00     71.0±1.12µs        ? ?/sec    1.01     71.6±1.19µs        ? ?/sec
map_1000_binary                          1.02      9.0±0.21ms        ? ?/sec    1.00      8.8±0.02ms        ? ?/sec
map_1000_binary_view                     1.00      8.8±0.07ms        ? ?/sec    1.01      8.8±0.26ms        ? ?/sec
map_1000_int32                           1.05      4.8±0.16ms        ? ?/sec    1.00      4.6±0.02ms        ? ?/sec
map_1000_utf8                            1.00      9.1±0.02ms        ? ?/sec    1.01      9.2±0.14ms        ? ?/sec
map_1000_utf8_view                       1.00      9.1±0.26ms        ? ?/sec    1.00      9.1±0.26ms        ? ?/sec
map_extract/int32/first/1024x32                                                 1.00     10.4±0.04µs        ? ?/sec
map_extract/int32/last/1024x1                                                   1.00      9.9±0.01µs        ? ?/sec
map_extract/int32/last/1024x32                                                  1.00     60.8±0.21µs        ? ?/sec
map_extract/int32/last/1x0                                                      1.00   733.6±84.35ns        ? ?/sec
map_extract/int32/last/1x1                                                      1.00  1103.2±61.80ns        ? ?/sec
map_extract/int32/missing/1024x32                                               1.00     57.6±0.05µs        ? ?/sec
map_extract/int32/varying/1024x32                                               1.00     36.9±0.15µs        ? ?/sec
map_extract/struct/first/1024x32                                                1.00     15.1±0.10µs        ? ?/sec
map_extract/struct/last/1024x1                                                  1.00     14.7±0.11µs        ? ?/sec
map_extract/struct/last/1024x32                                                 1.00    112.4±0.73µs        ? ?/sec
map_extract/struct/last/1x0                                                     1.00   628.8±74.98ns        ? ?/sec
map_extract/struct/last/1x1                                                     1.00  1028.2±57.90ns        ? ?/sec
map_extract/struct/missing/1024x32                                              1.00    107.2±0.94µs        ? ?/sec
map_extract/struct/varying/1024x32                                              1.00     67.1±0.41µs        ? ?/sec
map_extract/utf8_view/first/1024x32                                             1.00     15.8±0.04µs        ? ?/sec
map_extract/utf8_view/last/1024x1                                               1.00     15.8±0.03µs        ? ?/sec
map_extract/utf8_view/last/1024x32                                              1.00    205.4±0.73µs        ? ?/sec
map_extract/utf8_view/last/1x0                                                  1.00    940.9±6.59ns        ? ?/sec
map_extract/utf8_view/last/1x1                                                  1.00   1364.9±4.03ns        ? ?/sec
map_extract/utf8_view/missing/1024x32                                           1.00    188.3±0.23µs        ? ?/sec
map_extract/utf8_view/varying/1024x32                                           1.00    117.1±0.19µs        ? ?/sec

Resource Usage

map — base (merge-base)

Metric Value
Wall time 175.0s
Peak memory 132.3 MiB
Avg memory 38.3 MiB
CPU user 71.7s
CPU sys 0.1s
Peak spill 0 B

map — branch

Metric Value
Wall time 370.1s
Peak memory 139.3 MiB
Avg memory 93.1 MiB
CPU user 326.8s
CPU sys 0.3s
Peak spill 0 B

File an issue against this benchmark runner

@comphead comphead left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@neilconway

Copy link
Copy Markdown
Contributor Author

@comphead Thanks! That is a helpful suggestion. I poked around at the m[k] code and noticed at least one bug (#25082). Consolidating map_extract and m[k] likely makes sense, but it is involved enough that I'd prefer to do it as a separate PR. I filed #25083 for the consolidation work; I'll take that on once this lands.

@comphead comphead left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @neilconway

@neilconway

Copy link
Copy Markdown
Contributor Author

@jayzhan211 @comphead thanks for the reviews!

@neilconway
neilconway added this pull request to the merge queue Sep 8, 2026
@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to no response for status checks Sep 8, 2026
@neilconway
neilconway added this pull request to the merge queue Sep 8, 2026
Merged via the queue into apache:main with commit c29363f Sep 8, 2026
40 checks passed
@neilconway
neilconway deleted the neilc/fix-map-extract-comparator branch September 8, 2026 22:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation functions Changes to functions implementation sqllogictest SQL Logic Tests (.slt)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

map_extract fails to match equal struct keys with differently encoded dictionary nulls map_extract returns [NULL] instead of [] for missing keys

5 participants