feat: route map lookups through codegen dispatcher

docs/source/contributor-guide/benchmarking.md

@@ -43,6 +43,44 @@ These can also be run as a suite on a dedicated machine, see
 [Micro Benchmarking on AWS EC2](benchmarking_micro_ec2.md). Published results are in
 [benchmarks/results/micro](https://github.com/apache/datafusion-comet/tree/main/benchmarks/results/micro).
 
+## Map lookup dispatch
+
+`CometCodegenDispatchBenchmark --map-lookups` compares dispatcher on, dispatcher off
+(Spark projection over a Comet scan), and pure Spark. It checks answers and routes before
+timing `m[k]` and `element_at(m, k)`, with persisted DOUBLE/STRUCT keys, 4/64-entry maps,
+and lookup-only/mixed projections. Corpus generation is outside the measurements.
+
+Set `JAVA_HOME` to a JDK supported by the selected Spark profile. Build once, then reuse
+the release artifacts for separate benchmark JVMs. Forking Java also keeps Hadoop shutdown
+hooks outside Maven's disposable application classloader:
+
+```sh
+make release PROFILES=-Pspark-4.1
+comet_bench_opts=$(make -s print-benchmark-args BENCH_HEAP=4g PROFILES=-Pspark-4.1 | sed -n 's/^MAVEN_OPTS=//p')
+export COMET_CONF_DIR="$PWD/conf"
+cd spark
+../mvnw exec:exec -Pspark-4.1 -Dexec.classpathScope=test \
+  -Dexec.executable="$JAVA_HOME/bin/java" \
+  -Dexec.args="$comet_bench_opts -classpath %classpath org.apache.spark.sql.benchmark.CometCodegenDispatchBenchmark --map-lookups"
+```
+
+Use `--map-lookups --check-only` to validate the complete matrix without timing.
+For first use, replace `--map-lookups` at the end with
+`--map-lookups --case=get_map_value-double-4-lookup --first-use=dispatch`.
+Case IDs combine `get_map_value`/`element_at`, `double`/`struct`, `4`/`64`, and
+`lookup`/`mixed`. Repeat each selected case with `dispatch`, `fallback`, and `spark`
+in **separate JVMs**, before running another lookup in that JVM. The first two executions
+are timed before validation. Spark's compiled-source cache is shared across map sizes
+and projection shapes; resetting dispatcher counters does not clear it.
+
+Warmed tables use 1,024 and 65,536 rows, two seconds of warmup and at least two seconds of
+timing per arm, including a repeated fallback baseline to expose drift. Every query still
+creates new tasks: warmed does not mean task kernel setup is excluded. First-use results
+use 1,024 rows and report wall time, JVM-wide Spark compilation metrics, and task kernel
+initializations separately. First-minus-second is not an isolated compilation cost.
+Repeat measurements on an otherwise idle machine; these end-to-end scan/projection/sink
+results do not establish a universal speedup or isolate map traversal from bridge costs.
+
 ```{toctree}
 :hidden:
 

docs/source/contributor-guide/expression-audits/array_funcs.md

@@ -161,6 +161,7 @@
 - Spark 3.5.8 (audited 2026-05-27): baseline. `ElementAt(left, right, defaultValueOutOfBound, failOnError)`; group label `map_funcs`. Comet supports `ArrayType` input through native `ListExtract` and `MapType` input through native `map_extract`.
 - Spark 4.0.1 (audited 2026-05-27): `NullIntolerant` -> `nullIntolerant` field refactor; group label changes to `collection_funcs`; ANSI default flips to `true` so out-of-bound throws by default. Comet wires `failOnError` through to native `ListExtract`.
 - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1.
+- Dispatcher follow-up (audited 2026-09-11, all four versions): `CometElementAt` now mixes in `CodegenDispatchFallback` for map-key exclusions (see [map lookup audit](map_funcs.md#element_at)). This also routes its existing ANSI-mode nullable nondeterministic array/map exclusion through Spark's generated expression as a whole. The left operand is evaluated once and a NULL result skips the key/index, avoiding both duplicated stateful evaluation and eager index errors. Deterministic operands keep the native NULL guard; ordinary array lookup remains native. `element_at_ansi.sql` checks the nondeterministic array route and `element_at_map_ansi.sql` checks its map counterpart.
 
 ## flatten
 

docs/source/contributor-guide/expression-audits/map_funcs.md

@@ -23,10 +23,20 @@
 
 ## element_at
 
-- Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8.
-- Spark 3.5.8 (audited 2026-05-27): baseline. `ElementAt(left, right, defaultValueOutOfBound, failOnError) extends GetMapValueUtil`; the parser routes `element_at(<array>, ...)` to one overload and `element_at(<map>, ...)` to another. Comet routes `MapType` input through the same native `map_extract` path used by `GetMapValue`.
-- Spark 4.0.1 (audited 2026-05-27): adds `nullIntolerant: Boolean` field; semantics unchanged.
-- Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1.
+- Spark 3.4.3 (audited 2026-09-11): `ElementAt` uses `GetMapValueUtil` for map inputs. Interpreted lookup compares keys with `TypeUtils.getInterpretedOrdering`; generated lookup uses `CodegenContext.genEqual`. Both return the first matching value, or NULL for a missing key, NULL map/key or NULL value, independently of ANSI mode. The NULL left operand short-circuits key evaluation.
+- Spark 3.5.8 (audited 2026-09-11): same map evaluation and codegen; key-type checking uses `DataTypeUtils.sameType` instead of `DataType.sameType`.
+- Spark 4.0.1 (audited 2026-09-11): replaces the `NullIntolerant` trait with a field and refactors array error context; map lookup semantics are unchanged. `genEqual` now honors string collation, including inside arrays/structs. Map construction normalizes floating-point keys by default; lookup equality itself treats NaNs as equal and matches either sign of zero in all four versions.
+- Spark 4.1.1 (audited 2026-09-11): `ElementAt` class body is unchanged from 4.0.1.
+- Comet retains `MapKeySupport`'s native exclusions for floating-point keys at any nesting level, non-default collations and complex keys. `CodegenDispatchFallback` runs those lookups with Spark's generated code inside Comet, rather than relaxing native `map_extract` equality. ANSI-mode nullable nondeterministic map/array operands also dispatch as a whole, preserving single evaluation and NULL short-circuiting. Compatible deterministic operands retain the existing native route and ANSI NULL guard.
+- Regression coverage: `map_lookup_dispatch.sql` checks both lookup spellings on column inputs under ANSI on/off, including signed zero, NaN, infinities, subnormals, NULL/empty maps, NULL keys/values, nested NULLs, structural equality and NULL-map short-circuiting of a throwing key. `element_at_map_collation.sql` covers Spark 4.0+ collation; `CometMapExpressionSuite` retains constant-folding-on coverage and checks fallback when dispatch is disabled. The dispatcher type gate and expression enablement remain in control.
+
+## getmapvalue
+
+- Spark 3.4.3 (audited 2026-09-11): `GetMapValue` implements `map_col[key]` and shares the interpreted/generated lookup helpers with map `ElementAt`; see the equality and NULL behavior above.
+- Spark 3.5.8 (audited 2026-09-11): same `GetMapValue` implementation as 3.4.3.
+- Spark 4.0.1 (audited 2026-09-11): same lookup loop, with collation-aware equality supplied by code generation.
+- Spark 4.1.1 (audited 2026-09-11): tightens analysis-time input checking to `MapType`; runtime lookup is unchanged.
+- `CometMapExtract` uses `CodegenDispatchFallback` for the same key-type exclusions as `CometElementAt`. Native lookup and its guards are otherwise unchanged. SQL tests use map columns because Spark can rewrite a constructor subscript `map(...)[key]` into `CASE` even with constant folding disabled.
 
 ## map_contains_key
 

docs/source/user-guide/latest/expressions.md

@@ -165,7 +165,7 @@ The tables below list every Spark built-in expression with its current status.
 | `array_union` | ✅ | Native | NaN/signed-zero handling may differ ([details](compatibility/floating-point.md)) |
 | `arrays_overlap` | ✅ | Native |  |
 | `arrays_zip` | ✅ | Native |  |
-| `element_at` | ✅ | Native |  |
+| `element_at` | ✅ | Hybrid | Native array lookup; ANSI-mode nullable nondeterministic operands use the JVM codegen dispatcher |
 | `flatten` | ✅ | Native | Binary/struct/map elements fall back |
 | `get` | ✅ | — |  |
 | `sequence` | ✅ | Hybrid | Integral types run natively; date/timestamp sequences use codegen dispatch |
@@ -395,7 +395,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci
 
 | Function | Status | Implementation | Notes |
 | --- | --- | --- | --- |
-| `element_at` | ✅ | Native |  |
+| `element_at` | ✅ | Hybrid | Floating-point, non-default collated and complex map keys use the JVM codegen dispatcher; other supported keys remain native |
 | `map` | ✅ | Codegen dispatch | Routed through the JVM codegen dispatcher |
 | `map_concat` | ✅ | Codegen dispatch |  |
 | `map_contains_key` | ✅ | — |  |
@@ -407,6 +407,12 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci
 | `str_to_map` | ✅ | Hybrid |  |
 | `try_element_at` | ✅ | — | Lowers to `element_at` |
 
+The map subscript syntax `map_col[key]` (`GetMapValue`) uses the same hybrid key-type
+routing as `element_at`. Spark's generated lookup preserves NaN/signed-zero equality,
+collation rules and structural comparison of complex keys inside Comet. If codegen
+dispatch is disabled or cannot support the expression's types, the operator still
+falls back to Spark; the native key-type restrictions are unchanged.
+
 ---
 
 ## math_funcs

spark/src/main/scala/org/apache/comet/serde/arrays.scala

@@ -614,7 +614,7 @@ object CometArrayReverse extends CometExpressionSerde[Reverse] with ArraysBase {
 
 }
 
-object CometElementAt extends CometExpressionSerde[ElementAt] {
+object CometElementAt extends CometExpressionSerde[ElementAt] with CodegenDispatchFallback {
 
   /**
    * Under ANSI, neither native shape reproduces Spark for a nullable nondeterministic operand.
@@ -623,14 +623,22 @@ object CometElementAt extends CometExpressionSerde[ElementAt] {
    * whole batch first, so a throwing index fires on rows whose operand is NULL. `convert`
    * reproduces the short-circuit with a `CASE WHEN <operand> IS NOT NULL` guard, but that guard
    * serializes the operand twice, which a stateful operand cannot survive: the two copies advance
-   * its state independently and silently move values and NULLs. Declining leaves the lookup on
-   * Spark. Lifting this needs a native lookup that evaluates the operand once and masks the index
-   * evaluation with the result, at which point the guard becomes unnecessary for every operand.
+   * its state independently and silently move values and NULLs. Declining the native route sends
+   * the whole expression through Spark's codegen dispatcher, preserving single evaluation and
+   * NULL short-circuiting. A native implementation would need to evaluate the operand once and
+   * mask the index evaluation with the result, at which point the guard becomes unnecessary for
+   * every operand.
    */
   private val eagerIndexReason: String =
     "ANSI mode with a nullable nondeterministic array or map operand: a native lookup evaluates " +
       "the index over the whole batch, where Spark skips it on the rows whose operand is NULL"
 
+  private val inputTypeReason = "Input must be an array or map"
+  private val argumentsReason = "unsupported arguments for ElementAt"
+
+  override def getUnsupportedReasons(): Seq[String] =
+    MapKeySupport.unsupportedReasons ++ Seq(eagerIndexReason, inputTypeReason, argumentsReason)
+
   /** True when `convert` has to wrap the lookup to reproduce Spark's NULL short-circuit. */
   private def needsNullGuard(expr: ElementAt): Boolean =
     expr.failOnError && expr.left.nullable
@@ -642,7 +650,7 @@ object CometElementAt extends CometExpressionSerde[ElementAt] {
       expr.left.dataType match {
         case _: ArrayType => Compatible()
         case MapType(keyType, _, _) => MapKeySupport.keySupport(keyType)
-        case _ => Unsupported(Some("Input must be an array or map"))
+        case _ => Unsupported(Some(inputTypeReason))
       }
     }
   }
@@ -678,7 +686,7 @@ object CometElementAt extends CometExpressionSerde[ElementAt] {
               .setListExtract(arrayExtractBuilder)
               .build())
         } else {
-          withFallbackReason(expr, "unsupported arguments for ElementAt")
+          withFallbackReason(expr, argumentsReason)
           None
         }
     }

spark/src/main/scala/org/apache/comet/serde/maps.scala

@@ -46,6 +46,9 @@ private[serde] object MapKeySupport {
       "cannot reproduce Spark's equality for a complex key type (for example a `NULL` inside the " +
       "lookup key aborts the cast against a non-nullable nested component)."
 
+  val unsupportedReasons: Seq[String] =
+    Seq(floatingPointReason, collationReason, complexKeyReason)
+
   /**
    * The `SupportLevel` for a map-consuming expression whose stored-key type is `keyType`. Spark
    * finds a key with `TypeUtils.getInterpretedOrdering` over the keys `ArrayBasedMapBuilder`
@@ -114,7 +117,9 @@ object CometMapValues extends CometExpressionSerde[MapValues] {
   }
 }
 
-object CometMapExtract extends CometExpressionSerde[GetMapValue] {
+object CometMapExtract extends CometExpressionSerde[GetMapValue] with CodegenDispatchFallback {
+
+  override def getUnsupportedReasons(): Seq[String] = MapKeySupport.unsupportedReasons
 
   override def getSupportLevel(expr: GetMapValue): SupportLevel = expr.child.dataType match {
     case MapType(keyType, _, _) => MapKeySupport.keySupport(keyType)

spark/src/test/resources/sql-tests/expressions/array/element_at_ansi.sql

@@ -101,14 +101,15 @@ query
 SELECT id, element_at(CASE WHEN id <> 2 THEN array(1) END, 1) AS v
 FROM ansi_element_at_null
 
--- A nondeterministic operand is declined outright. Neither native shape reproduces Spark: the
+-- A nondeterministic operand is dispatched as a whole. Neither native shape reproduces Spark: the
 -- `CASE WHEN <array> IS NOT NULL` guard serializes the operand twice, so a stateful operand's two
 -- copies drift, and the unguarded lookup evaluates the index over the whole batch, raising
--- DIVIDE_BY_ZERO at id = 2 on the very row whose array is NULL. `rand(7L) < 2` is always true, so
--- both operands are NULL on every row and Spark returns NULL without evaluating either index.
+-- DIVIDE_BY_ZERO at id = 2 on the very row whose array is NULL. The first operand alternates
+-- NULL/non-NULL by row. `rand(7L) < 2` is always true, so the second operand is always NULL and
+-- Spark returns NULL without evaluating its throwing index.
 -- The non-ANSI spelling stays native and is covered in element_at.sql.
 -- https://github.com/apache/datafusion-comet/issues/5544
-query expect_fallback(nullable nondeterministic array or map operand)
+query expect_dispatch(element_at)
 SELECT id,
        element_at(IF(monotonically_increasing_id() % 2 = 0, CAST(NULL AS ARRAY<INT>), array(1)), 1) AS v1,
        element_at(IF(rand(7L) < 2, CAST(NULL AS ARRAY<INT>), array(1)), 1 + (id % (id - 2))) AS v2

spark/src/test/resources/sql-tests/expressions/map/element_at_map.sql

@@ -25,7 +25,7 @@ INSERT INTO test_element_at_map VALUES
   (NULL, NULL)
 
 -- key found
-query
+query expect_native(element_at)
 SELECT element_at(m, 'a'), element_at(m, 'b') FROM test_element_at_map
 
 -- key not found → NULL
@@ -52,29 +52,29 @@ SELECT element_at(mi, CAST(1 AS BIGINT)), element_at(mi, CAST(2 AS SMALLINT)) FR
 query
 SELECT element_at(map('a', 1, 'b', 2), 'a'), element_at(map('a', 1, 'b', 2), 'missing'), element_at(map('a', 1, 'b', 2), NULL)
 
--- Map key types whose Spark equality Comet's native `map_extract` cannot reproduce fall back to
--- Spark. These stay on the constructor path here because the SQL harness excludes
+-- Map key types whose Spark equality Comet's native `map_extract` cannot reproduce dispatch to
+-- Spark's generated code. These stay on the constructor path here because the SQL harness excludes
 -- `ConstantFolding`; `CometMapExpressionSuite` covers the folded-literal form of each.
 
--- Spark stores `-0.0` map keys as `+0.0` and compares with `nanSafeCompareDoubles`, so a `-0.0`
--- lookup finds the `+0.0` key. Native lookup compares the raw Arrow values.
-query expect_fallback(Spark normalizes floating-point map keys)
-SELECT element_at(map(CAST(0 AS DOUBLE), 7), CAST(-0.0 AS DOUBLE))
+-- Spark's floating-point equality treats both signs of zero as equal, so a `-0.0` lookup finds
+-- the `+0.0` key. Native lookup compares the raw Arrow values.
+query expect_dispatch(element_at)
+SELECT element_at(map(CAST(0 AS DOUBLE), 7), double('-0.0'))
 
-query expect_fallback(Spark normalizes floating-point map keys)
-SELECT element_at(map(CAST(0 AS FLOAT), 7), CAST(-0.0 AS FLOAT))
+query expect_dispatch(element_at)
+SELECT element_at(map(CAST(0 AS FLOAT), 7), float('-0.0'))
 
 -- The floating-point decline walks every nesting level of the key type, so an array-of-double key
--- falls back for the same reason.
-query expect_fallback(Spark normalizes floating-point map keys)
-SELECT element_at(map(array(CAST(0 AS DOUBLE)), 7), array(CAST(-0.0 AS DOUBLE)))
+-- dispatches for the same reason.
+query expect_dispatch(element_at)
+SELECT element_at(map(array(CAST(0 AS DOUBLE)), 7), array(double('-0.0')))
 
 -- A complex key type: `map_extract` casts the lookup key to the map's exact Arrow key type, so a
 -- NULL inside the lookup key would abort the cast instead of missing the lookup.
-query expect_fallback(casts the lookup key to the map's exact Arrow key type)
+query expect_dispatch(element_at)
 SELECT element_at(map(array(1), 7), array(CAST(NULL AS INT)))
 
-query expect_fallback(casts the lookup key to the map's exact Arrow key type)
+query expect_dispatch(element_at)
 SELECT element_at(map(named_struct('a', 1), 7), named_struct('a', 1))
 
 -- `BinaryType` keys need no decline: Arrow compares them by content, as Spark's ordering does.

spark/src/test/resources/sql-tests/expressions/map/element_at_map_ansi.sql

@@ -34,3 +34,11 @@ INSERT INTO test_element_at_nested_ansi VALUES (1), (2), (3)
 query
 SELECT id, element_at(element_at(map(1, map(0, 7)), id), id % (id - 2)) AS v
 FROM test_element_at_nested_ansi
+
+-- The nullable nondeterministic operand must be evaluated once inside the dispatcher. In
+-- particular, the second lookup must not evaluate its throwing key on the NULL-map row.
+query expect_dispatch(element_at)
+SELECT id,
+       element_at(IF(monotonically_increasing_id() % 2 = 0, CAST(NULL AS MAP<INT, INT>), map(1, 7)), 1) AS v1,
+       element_at(IF(rand(7L) < 2, CAST(NULL AS MAP<INT, INT>), map(1, 7)), id % (id - 2)) AS v2
+FROM test_element_at_nested_ansi

spark/src/test/resources/sql-tests/expressions/map/element_at_map_collation.sql

@@ -19,8 +19,8 @@
 
 -- Spark 4.0+ supports string collations. Spark compares a map key under its declared collation,
 -- so a `UTF8_LCASE` map matches `A1` against a dynamic `a1` lookup and returns `7`. Comet's native
--- `map_extract` compares string keys as `UTF8_BINARY`, so `CometElementAt` declines the lookup and
--- Spark evaluates the projection. The `element_at` form is used rather than `m[key]` because
+-- `map_extract` compares string keys as `UTF8_BINARY`, so `CometElementAt` dispatches the lookup to
+-- Spark's generated code inside Comet. The `element_at` form is used rather than `m[key]` because
 -- Spark's `SimplifyExtractValueOps` rewrites `map(...)[key]` over a literal map into a `CASE`
 -- before it can reach the native map lookup.
 
@@ -30,6 +30,23 @@ CREATE TABLE test_element_at_collation(k string) USING parquet
 statement
 INSERT INTO test_element_at_collation VALUES ('a1'), ('A1'), ('zz'), (NULL)
 
-query expect_fallback(cannot honour a non-default collation)
+query expect_dispatch(element_at)
 SELECT element_at(map(CAST('A1' AS STRING COLLATE UTF8_LCASE), 7), CAST(k AS STRING COLLATE UTF8_LCASE))
 FROM test_element_at_collation
+
+-- Column inputs keep GetMapValue intact, rather than rewriting map(...)[key] into CASE.
+statement
+CREATE TABLE test_lookup_collation(m MAP<STRING, INT>, k STRING) USING parquet
+
+statement
+INSERT INTO test_lookup_collation VALUES
+  (map('A1', 7), 'a1'), (map('A1', 8), 'zz'), (map('A1', NULL), 'a1'),
+  (map(), 'a1'), (NULL, 'a1'), (map('A1', 9), NULL)
+
+query expect_dispatch(getmapvalue)
+SELECT CAST(m AS MAP<STRING COLLATE UTF8_LCASE, INT>)[CAST(k AS STRING COLLATE UTF8_LCASE)]
+FROM test_lookup_collation
+
+query expect_dispatch(element_at)
+SELECT element_at(CAST(m AS MAP<STRING COLLATE UTF8_LCASE, INT>), CAST(k AS STRING COLLATE UTF8_LCASE))
+FROM test_lookup_collation