andygrove opened a new pull request, #5806:
URL: https://github.com/apache/datafusion-comet/pull/5806
## Which issue does this PR close?
Closes #5795.
## Rationale for this change
`GetMapValue` (`m[k]`) and `element_at(<map>, k)` both serialise to the
DataFusion `map_extract`
UDF, and the planner then unwrapped its one-element list result with a
second `ListExtract` pass.
The diagnosis in the issue holds. `general_map_extract_inner` re-slices the
query key and every
candidate key into a fresh `ArrayRef` per comparison and compares them
through `dyn Array`
equality:
```rust
let query_key = query_keys_array.slice(row_index, 1);
let value_index = (0..len).find(|&i| keys.slice(start + i, 1).as_ref() ==
query_key.as_ref());
```
so the lookup is O(rows x entries-per-row) allocations, and the constant key
is not hoisted out of
the row loop. That is why the map lookup was the one map operation where
Comet lost to Spark.
## What changes are included in this PR?
A native `SparkMapExtract`, registered under the same `map_extract` name so
it overrides the
DataFusion one, in `native/spark-expr/src/map_funcs/map_extract.rs`:
- one Arrow `eq` over the whole batch of map entries (broadcasting the key
when it is a constant),
then a bitmask scan for each row's first match, then one `take` to gather
the values;
- it returns the matched value directly instead of a one-element list, so
the `ListExtract` wrapper
in `planner.rs` is no longer needed and that second pass goes away;
- a sliced `MapArray` keeps its original entry offsets, so only the visible
entry window is
compared.
Semantics are unchanged: the first matching entry wins, and a missing key, a
`NULL` map row and a
`NULL` lookup key all yield `NULL`. Key types whose Spark equality a native
lookup cannot reproduce
(floating point, non-default collations, complex keys) are still declined by
`MapKeySupport`, so
they never reach the kernel. Arrow's `eq` rejects nested key types, so
DataFusion's element-wise
comparison is kept as a backstop for those, and a lookup key whose runtime
type is not the map's
key type is now rejected rather than silently missing every row.
## How are these changes tested?
**Existing tests.** `CometMapExpressionSuite` (25),
`CometArrayExpressionSuite` (64),
`CometExpressionSuite` (141) and the full `CometSqlFileTestSuite` (474,
including
`element_at_map.sql`, `element_at_map_ansi.sql`,
`element_at_map_collation.sql` and
`get_map_value.sql`) all pass.
**New tests.** 13 Rust unit tests covering hit/miss, duplicate keys, `NULL`
maps, `NULL` values,
`NULL` and per-row lookup keys, sliced maps, empty input, a scalar map
argument, non-string keys,
the nested-key backstop, and the rejected argument types. Two new
`CometMapExpressionSuite` cases
cover the same two tricky paths end to end: `element_at` under a native
`OFFSET` (sliced map) and
`element_at`/`m[k]` with a per-row lookup key.
**Kernel benchmark** (new `native/spark-expr/benches/map_extract.rs`,
8192-row batches of
`map<string, string>`, 10% `NULL` rows, 60 distinct keys, run against both
implementations):
| entries/map | case | `datafusion` | `comet` | speedup |
| --- | --- | --- | --- | --- |
| 2 | constant key | 1.873 ms | 41.8 us | 44.9x |
| 8 | constant key | 6.181 ms | 110.4 us | 56.0x |
| 32 | constant key | 18.85 ms | 376.8 us | 50.0x |
| 2 | per-row key | 1.946 ms | 145.0 us | 13.4x |
| 8 | per-row key | 6.390 ms | 470.7 us | 13.6x |
| 32 | per-row key | 19.68 ms | 1.795 ms | 11.0x |
**End to end**, reproducing the issue's shape (2,000,000 rows, `attrs
map<string, string>`, 10%
`NULL` maps, 0-6 entries each, 60 distinct keys, 400 distinct values; `noop`
sink, median of 5
after 2 warmups; Apple silicon, so the absolute numbers differ from the
issue's Ryzen run):
| query | Spark | Comet before | Comet after |
| --- | --- | --- | --- |
| `SELECT attrs FROM t` | 1098 ms | 716 ms | 856 ms |
| `SELECT size(attrs) FROM t` | 291 ms | 214 ms | 210 ms |
| `SELECT size(map_keys(attrs)) FROM t` | 289 ms | 246 ms | 214 ms |
| `SELECT size(map_values(attrs)) FROM t` | 276 ms | 231 ms | 204 ms |
| `SELECT element_at(attrs, 'a1') FROM t` | 506 ms | 278 ms | **211 ms** |
| combined query from the issue | 602 ms | 307 ms | 273 ms |
Taking `size(attrs)` as the control, as the issue does, the isolated cost of
the lookup goes from
**+64 ms to +1.5 ms**, and `element_at` is no longer the outlier among the
map kernels. The
Spark-arm numbers drifted about 15% between the two runs on this machine, so
read the Comet columns
against each other and against their own control rather than the ratios.
--
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]
