perf: Unify, optimize map[k] and map_extract(key) - #25201
neilconway wants to merge 5 commits into
Conversation
github-actions
Bot
added
sqllogictest
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.91% 81.91%
========================================
Files 1134 1134
Lines 425631 425769 +138
Branches 425631 425769 +138
========================================
+ Hits 348647 348775 +128
- Misses 56304 56306 +2
- Partials 20680 20688 +8 ☔ 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
56f7828 to
b33a1c0
Compare
neilconway
changed the title
map[k] and map_extract(key)map[k] and map_extract(key)
|
Nice, I'll check it tomorrow, and might be related to apache/datafusion-comet#5806 |
comphead
left a comment
comphead
left a comment
There was a problem hiding this comment.
Reviewed with the focus on the new map_lookup kernel. The consolidation is the right call, the NULLIF null-row fix is a genuine bug fix, and the benchmark matrix is unusually thorough.
One behavioral contract is broken; the rest are coverage notes. Inline below.
- P2 the match hint can return a non-first entry for rows with duplicate keys, which makes the result depend on row order and therefore on
batch_size, and contradictsmap_lookup's own rustdoc. Confirmed with a failing test. - P2 no test reaches the vectorized
eqbranch on a sliced map, and nothing asserts the two scan strategies agree. - P3 an untyped
NULLkey still errors while a typedNULLkey now returnsNULL. - P3 the
get_fieldbench dropped the 8192-row shape, which is the defaultbatch_size. - P3
map_extractgot looser for nested keys: the old exactDataTypeequality check is now justis_nested(), andmake_comparatorcompares struct fields positionally while ignoring names, somap_extract(map<struct<a:int>,v>, struct<b:int>)now matches instead of erroring. Intentional and documented in the rustdoc, but worth a line in the PR description since it is user-visible for directinvoke_with_argscallers.
Fuzz suggestion: map_lookup is a pure function with an obvious oracle (naive per-row find). A proptest over row count, entries per row, null rows, null keys, slice offsets, duplicate keys and key types could assert three properties: comparator result == eq result, batch-size invariance (chunk and concatenate == whole), and slice invariance.
| let hinted = start + self.hint; | ||
| let found = if hinted < end && is_match(hinted, row) { | ||
| Some(hinted) | ||
| } else { | ||
| (start..end).find(|&entry| entry != hinted && is_match(entry, row)) | ||
| }; |
There was a problem hiding this comment.
P2: this can return a later duplicate rather than the first match, contradicting the rustdoc above ("the first entry whose key equals that row's lookup key"). self.hint carries over from the previous row, so for a row that holds the lookup key twice the hint can land on the second one.
Repro, added to the map_lookup_tests module on this branch:
// keys [1, 7 | 7, 7], two rows of length 2, lookup key 7
let got = map_lookup(&map, &Int32Array::from(vec![7])).unwrap();
// first match per row would be [1, 2]; actual is [1, 3]Reordering so row 0 matches at position 0 (keys [7, 1 | 7, 7]) gives [0, 2], the first match. So the answer depends on the preceding row, and therefore on batch_size, since the hint is per call and each call is one batch. Both pre-PR implementations used (start..end).find(...) and always returned the first match.
Arrow does not enforce key uniqueness (MapArray::try_new validates offsets, nulls and child count only), so duplicates arrive through Parquet, IPC, FFI and custom TableProviders. DataFusion's own map() / make_map() reject them, so this is not reachable from a SQL literal.
Either scan start..hinted when the hint matches, or state in the rustdoc that the entry returned for a row with duplicate keys is unspecified. The doc as written is the part that is wrong today.
| let range_start = offsets[sample] as usize; | ||
| let in_range = map_keys.slice(range_start, last - range_start); | ||
| let matches = eq(&Scalar::new(keys.slice(0, 1)), &in_range)?; | ||
| // Neither side has nulls, so the value bits alone are meaningful. | ||
| let bits = matches.values(); | ||
| scanner.scan(rest, |entry, _| bits.value(entry - range_start)); |
There was a problem hiding this comment.
Test gap: this branch is never reached on a sliced map. sliced_map_and_keys has 4 rows, so rest is empty and the branch is skipped; vectorized_scan_after_missing_sample has offsets[0] == 0, so range_start is only ever exercised against a zero base.
I verified the indexing is correct today (200-row rotated-key map, slice(37, 150), matches a naive per-row scan), so this is coverage rather than a live bug. But sliced map batches with more than 32 rows are common after limit, coalesce and batch splitting, and an off-by-offsets[0] here silently returns values from the wrong entries instead of erroring.
| if single_key | ||
| && !key_type.is_nested() | ||
| && !rest.is_empty() | ||
| && comparisons * 2 > sampled_entries |
There was a problem hiding this comment.
Nothing asserts the two strategies agree. The rustdoc makes a specific semantic claim (total ordering, so -0.0 and 0.0 are different keys and NaN matches NaN), but map_extract_float_keys uses a single-row map and can therefore only reach the comparator.
They agree today only because ArrowNativeTypeOp::is_eq for floats is bitwise and eq routes through it, which is an arrow implementation detail rather than a stable contract. If it moves to IEEE equality, map[0.0] starts matching a -0.0 key only for batches larger than 32 rows where this heuristic flips, and nothing here catches it. SAMPLE_ROWS is private and the threshold is inline, so there is no seam to force a strategy in a test; a test-only override would make the equivalence test easy to write.
Separately on the threshold itself: the description says the measured break-even is around 60% but this uses 50%, and map_extract/utf8_view/shuffled/1024x32 (+19.8%) is exactly the marginal case. At width 32 with a rotated key the sample compares about 16 entries per row plus the failed hint probe, so comparisons * 2 lands just over sampled_entries and flips to eq. Moving toward the measured value may recover that case without giving up the missing wins.
| let compatible = if key_type.is_nested() { | ||
| keys.data_type().is_nested() | ||
| } else { | ||
| strip_dictionary(key_type).equals_datatype(strip_dictionary(keys.data_type())) | ||
| }; | ||
| if !compatible { | ||
| return exec_err!( |
There was a problem hiding this comment.
DataType::Null keys fail here, before reaching the null short-circuit below, so SELECT column1[NULL] still errors (the .slt expectation updated in this PR) while the new MAP {1:'a', 2:'b'}[arrow_cast(NULL, 'Int64')] case returns NULL. Two spellings of the same SQL NULL behave differently.
Not a regression, but this PR rewrites the path and changes the message, so it is the natural place to fix it: accept DataType::Null and return UInt32Array::new_null(map.len()).
| /// Every tenth row is null. | ||
| fn map_array(size: usize, entries: usize) -> ArrayRef { | ||
| let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); | ||
| const ROWS: usize = 1024; |
There was a problem hiding this comment.
The old bench had an 8192-row case and this pins ROWS = 1024. 8192 is DataFusion's default batch_size, and the design amortizes a fixed 32-row comparator sample over the batch, so 1024 rows both understates the benefit and overstates the sample cost. Worth keeping one 8192-row shape, plus something just above the threshold (say 40 rows, as produced by a selective filter) where the sample is a large fraction of the batch rather than 0.4% of it.
Also, swapping plain Utf8 for Utf8View drops a common shape: map<string, ...> read without schema_force_view_types takes a different comparator (compare_bytes vs compare_byte_view).
3 participants
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.