neilconway opened a new pull request, #25201:
URL: https://github.com/apache/datafusion/pull/25201
## Which issue does this PR close?
- Closes ##25083.
## Rationale for this change
`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).
We can also use `take` to construct the results, which is faster than the
previous approach based on `MutableArrayData::extend`.
Benchmarks: (M4 Max)
get_field (map[key]):
- get_field_map/utf8/first/1024x4: 15.60 µs -> 5.44 µs, -65.2%
- get_field_map/utf8/last/1024x4: 17.01 µs -> 5.33 µs, -68.7%
- get_field_map/utf8/missing/1024x4: 15.91 µs -> 13.02 µs, -18.2%
- get_field_map/utf8/first/1024x32: 33.79 µs -> 5.59 µs, -83.5%
- get_field_map/utf8/last/1024x32: 71.68 µs -> 6.29 µs, -91.2%
- get_field_map/utf8/missing/1024x32: 70.74 µs -> 77.22 µs, +9.2%
- get_field_map/utf8/first/8192x4: 123.02 µs -> 40.90 µs, -66.8%
- get_field_map/utf8/last/8192x4: 134.52 µs -> 40.27 µs, -70.1%
- get_field_map/utf8/missing/8192x4: 120.77 µs -> 101.46 µs, -16.0%
- get_field_map/int32/first/1024x4: 6.73 µs -> 3.25 µs, -51.7%
- get_field_map/int32/last/1024x4: 8.62 µs -> 3.26 µs, -62.1%
- get_field_map/int32/missing/1024x4: 7.86 µs -> 4.61 µs, -41.4%
- get_field_map/int32/first/1024x32: 9.71 µs -> 3.48 µs, -64.1%
- get_field_map/int32/last/1024x32: 21.65 µs -> 3.58 µs, -83.4%
- get_field_map/int32/missing/1024x32: 19.72 µs -> 19.96 µs, +1.2%
- get_field_map/struct/first/1024x4: 8.24 µs -> 3.80 µs, -53.9%
- get_field_map/struct/last/1024x4: 13.51 µs -> 3.82 µs, -71.7%
- get_field_map/struct/missing/1024x4: 13.13 µs -> 10.53 µs, -19.8%
- get_field_map/struct/first/1024x32: 8.44 µs -> 4.06 µs, -51.9%
- get_field_map/struct/last/1024x32: 65.12 µs -> 4.21 µs, -93.5%
- get_field_map/struct/missing/1024x32: 60.52 µs -> 58.88 µs, -2.7%
map_extract:
- map_extract/int32/last/1x0: 286 ns -> 376 ns, +31.5%
- map_extract/utf8_view/last/1x0: 362 ns -> 462 ns, +27.8%
- map_extract/struct/last/1x0: 239 ns -> 322 ns, +34.8%
- map_extract/int32/last/1x1: 395 ns -> 371 ns, -6.1%
- map_extract/utf8_view/last/1x1: 502 ns -> 499 ns, -0.6%
- map_extract/struct/last/1x1: 396 ns -> 388 ns, -1.9%
- map_extract/int32/last/1024x1: 6.22 µs -> 3.19 µs, -48.7%
- map_extract/utf8_view/last/1024x1: 8.66 µs -> 5.04 µs, -41.8%
- map_extract/struct/last/1024x1: 8.68 µs -> 4.52 µs, -47.9%
- map_extract/int32/first/1024x32: 6.36 µs -> 4.29 µs, -32.6%
- map_extract/int32/last/1024x32: 38.82 µs -> 4.44 µs, -88.6%
- map_extract/int32/missing/1024x32: 35.42 µs -> 21.62 µs, -39.0%
- map_extract/int32/varying/1024x32: 25.39 µs -> 24.59 µs, -3.2%
- map_extract/utf8_view/first/1024x32: 10.19 µs -> 7.00 µs, -31.4%
- map_extract/utf8_view/last/1024x32: 122.13 µs -> 7.21 µs, -94.1%
- map_extract/utf8_view/missing/1024x32: 113.76 µs -> 79.67 µs, -30.0%
- map_extract/utf8_view/varying/1024x32: 71.81 µs -> 66.93 µs, -6.8%
- map_extract/struct/first/1024x32: 8.81 µs -> 4.85 µs, -45.0%
- map_extract/struct/last/1024x32: 71.39 µs -> 5.02 µs, -93.0%
- map_extract/struct/missing/1024x32: 63.47 µs -> 63.91 µs, +0.7%
- map_extract/struct/varying/1024x32: 41.39 µs -> 38.58 µs, -6.8%
## What changes are included in this PR?
* Add `datafusion_functions::utils::map_lookup`, which returns for each map
row the index of the first matching entry or null.
* Implement `get_field` and `map_extract` on top of the shared `map_lookup`
helper, constructing the results with `take`
* In `map_extract`, optimize for the single-scalar-key case by passing it
through as a scalar value instead of expanding it to the batch size
* Extend benchmarks
## What is the testing strategy for this PR?
Existing tests pass; new tests added.
## Are there any user-facing changes?
No, aside from some corner-case changes like how error messages are
formatted.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
