perf: Unify, optimize map[k] and map_extract(key) - #25201
Open
neilconway wants to merge 4 commits into
Open
Conversation
github-actions
Bot
added
sqllogictest
Contributor
Author
Both suites cover Int32, Utf8View, and struct keys on maps of 4 and 32
entries, named `{key type}/{lookup}/{rows}x{entries}`. The lookups are the
first and last entry of every row, a key that every row holds at a
different position (`shuffled`), and a key present in no row, plus one key
per row for `map_extract`. Cases that only rescaled another case, such as
8192 rows or one-entry maps, are dropped; `map_extract` keeps two
single-row cases that measure per-batch fixed cost.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #25201 +/- ##
========================================
Coverage 81.93% 81.93%
========================================
Files 1133 1133
Lines 423529 423667 +138
Branches 423529 423667 +138
========================================
+ Hits 347032 347150 +118
- Misses 55907 55918 +11
- Partials 20590 20599 +9 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
`map[key]` (aka `get_field`) and `map_extract(map, key)` collectively
had three ways to lookup keys in a map:
(1): `get_field` used a per-batch `eq` kernel for scalar keys. This is
efficient for maps with many entries where `key` is not found, but
slower for maps in which the key can be found quickly (because `eq`
does not allow early-stopping). On my local machine, `eq` only
beats a comparator-based approach if the latter required touching
more than 60% of the keys in a row.
(2): `get_field` used a comparator-based approach for nested keys.
(3): `map_extract` used a comparator-based approach for all keys.
Conceptually, these two functions only differ in how the result is
represented, so it makes sense to consolidate them. We can also adopt a
hybrid strategy that gets the best of the previous approaches for most
inputs:
* Start with a comparator-based approach.
* After the first row, remember the index at which the matching key was
found, and check that index first for subsequent rows. This takes
advantage of the observation that most map rows have their keys in
the same order.
* After 32 rows, check whether the comparator looked at more than 75% of
the entries in those rows. If it did, the early stopping that the
comparator approach allows is not useful and we switch to an `eq`
kernel for the remainder of the batch (as long as the map key is not a
nested type).
Add `datafusion_functions::utils::map_lookup`, which returns for each map
row the index of the first matching entry or null, and build both
functions on it with `take`. The lookup scans each row with a comparator,
which stops at the first match, and tries the position where the previous
row matched first, since rows in a batch usually share key order; a hit
then usually costs one comparison wherever the key sits. A miss has no
match to stop at, so every entry of the row is compared at the
comparator's higher per-comparison cost. After sampling 32 rows, if most
of them missed, the remaining entries are compared with one vectorized
`eq` instead, which does the same full comparison more cheaply. A lookup
key that differs from the map key type only in dictionary encoding is cast
rather than rejected.
`map_extract` now passes a scalar key through as a single row instead of
expanding it to the batch size. `map[key]` reports a mismatched key type
as an execution error instead of an Arrow comparison error, and a null map
row that still carries entries yields NULL, as it already did from
`map_extract`.
Closes apache#25083
neilconway
force-pushed
the
neilc/map-lookup-kernel
branch
from
September 11, 2026 21:33
56f7828 to
b33a1c0
Compare
neilconway
changed the title
map[k] and map_extract(key)map[k] and map_extract(key)
Contributor
|
Nice, I'll check it tomorrow, and might be related to apache/datafusion-comet#5806 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
get_fieldandmap_extract#25083.map[key]misbehaves for NULL map rows that have entries #25210.Rationale for this change
map[key](akaget_field) andmap_extract(map, key)collectivelyhad three ways to lookup keys in a map:
(1):
get_fieldused a per-batcheqkernel for scalar keys. This isefficient for maps with many entries where
keyis not found, butslower for maps in which the key can be found quickly (because
eqdoes not allow early-stopping). On my local machine,
eqonlybeats a comparator-based approach if the latter required touching
more than 60% of the keys in a row.
(2):
get_fieldused a comparator-based approach for nested keys.(3):
map_extractalways used a comparator-based approach.Conceptually, these two functions only differ in how the result is
represented, so it makes sense to consolidate them. We can also adopt a
hybrid strategy that gets the best of the previous approaches for most
inputs:
found, and check that index first for subsequent rows. This takes
advantage of the observation that most map rows have their keys in
the same order.
the entries in those rows. If it did, the early stopping that the
comparator approach allows is not useful and we switch to an
eqkernel for the remainder of the batch (as long as the map key is not a
nested type).
We can also use
taketo construct the results, which is faster than the previous approach based onMutableArrayData::extend.Along the way, fix
get_fieldto behave correctly for NULL map rows with non-empty offset ranges (#25210).Benchmarks: (M4 Max)
get_field (map[key]):
map_extract:
What changes are included in this PR?
datafusion_functions::utils::map_lookup; for each map row, this returns the index of the first matching entry or null.get_fieldandmap_extracton top of the sharedmap_lookuphelper, constructing the results withtakemap_extract, optimize for the single-scalar-key case by passing it through as a scalar value instead of expanding it to the batch sizemap[key]misbehaves for NULL map rows that have entries #25210What is the testing strategy for this PR?
Existing tests pass; new tests added.
Are there any user-facing changes?
No, aside from the bugfix and some corner-case changes like how error messages are formatted.