apache / apache/datafusion-comet

Map lookups with float, collated or complex keys fall back to Spark (`map_col[key]`, `element_at`)

Open
#5,580 1 comment 0 reactions 1 assignee Claimed by @LinSimon-901101 View on GitHub
area:expressions enhancement map expressions
Dominant language
Scala
Stars
1.3k
Forks
373
Avg merge
2d 4h
Merged PRs (30d)
198

Description

Part of #5572.

`MapKeySupport.keySupport` (`spark/src/main/scala/org/apache/comet/serde/maps.scala:68`) returns `Unsupported` for three classes of map key:

- floating point anywhere in the key type, because Spark normalizes `-0.0` to `+0.0` and canonicalizes `NaN` in `ArrayBasedMapBuilder` while native `map_extract` compares raw Arrow values
- a non-default collation, because native compares as `UTF8_BINARY`
- any complex key type, because `map_extract`'s coercion casts the lookup key to the map's exact Arrow key type

Two serdes consume it, and neither mixes in `CodegenDispatchFallback`:

- `CometMapExtract` (`GetMapValue`, i.e. `map_col[key]`) — `serde/maps.scala:117`, gate at `:120`
- `CometElementAt` on map input — `serde/arrays.scala:614`

So `map_col[key]` and `element_at(map_col, key)` fail the whole projection back to Spark whenever the map has double, collated-string or struct keys.

The analysis behind each decline is right — these really are cases the native lookup gets wrong — which is exactly why the dispatcher is the correct answer rather than a native fix. Spark's own `GetMapValue.doGenCode` and `ElementAt.doGenCode` give the normalization and interpreted-ordering equality for free, and `MapType` is fully supported by `CometBatchKernelCodegen.isSupportedDataType` (recursively, as long as the leaves are).

Fix: mix `CodegenDispatchFallback` into both serdes.

Note that `CometElementAt`'s other `Unsupported` arm, "Input must be an array or map", is unreachable — Spark's own type checking rejects that before Comet sees it — so it needs no handling.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.