You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/01 04:40:57 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-5875-92b43cfaaea476ad8c146f9a252c640594921766 created (now 58a540a9ed)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-5875-92b43cfaaea476ad8c146f9a252c640594921766
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at 58a540a9ed feat: route map lookups through codegen dispatcher (#5875)
This branch includes the following new commits:
new 58a540a9ed feat: route map lookups through codegen dispatcher (#5875)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: feat: route map lookups through codegen dispatcher (#5875)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-5875-92b43cfaaea476ad8c146f9a252c640594921766
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 58a540a9ed4c948f2e13cd248c2c01a25ac5e5f1
Author: LinSimon-901101 <bc...@gm.ntcu.edu.tw>
AuthorDate: Thu Oct 1 04:34:09 2026 +0000
feat: route map lookups through codegen dispatcher (#5875)
* feat: route map lookups through codegen dispatcher
* bench: add map lookup codegen dispatch benchmarks
* test: tighten map lookup dispatch assertions and documentation
---
docs/source/contributor-guide/benchmarking.md | 38 +++
.../expression-audits/array_funcs.md | 1 +
.../expression-audits/map_funcs.md | 18 +-
docs/source/user-guide/latest/expressions.md | 10 +-
.../main/scala/org/apache/comet/serde/arrays.scala | 20 +-
.../main/scala/org/apache/comet/serde/maps.scala | 7 +-
.../expressions/array/element_at_ansi.sql | 9 +-
.../sql-tests/expressions/map/element_at_map.sql | 28 +-
.../expressions/map/element_at_map_ansi.sql | 8 +
.../expressions/map/element_at_map_collation.sql | 23 +-
.../sql-tests/expressions/map/get_map_value.sql | 12 +-
.../expressions/map/map_lookup_dispatch.sql | 125 +++++++++
.../org/apache/comet/CometMapExpressionSuite.scala | 75 +++--
.../benchmark/CometCodegenDispatchBenchmark.scala | 303 +++++++++++++++++++--
14 files changed, 592 insertions(+), 85 deletions(-)
diff --git a/docs/source/contributor-guide/benchmarking.md b/docs/source/contributor-guide/benchmarking.md
index c3318515e4..887d0fb7d1 100644
--- a/docs/source/contributor-guide/benchmarking.md
+++ b/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:
diff --git a/docs/source/contributor-guide/expression-audits/array_funcs.md b/docs/source/contributor-guide/expression-audits/array_funcs.md
index cd167885d0..9993d9f0f1 100644
--- a/docs/source/contributor-guide/expression-audits/array_funcs.md
+++ b/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 operan [...]
## flatten
diff --git a/docs/source/contributor-guide/expression-audits/map_funcs.md b/docs/source/contributor-guide/expression-audits/map_funcs.md
index f50475a449..16b28a959f 100644
--- a/docs/source/contributor-guide/expression-audits/map_funcs.md
+++ b/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 N [...]
+- 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 expressio [...]
+
+## 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
diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md
index 1487e4634e..8a304954c8 100644
--- a/docs/source/user-guide/latest/expressions.md
+++ b/docs/source/user-guide/latest/expressions.md
@@ -170,7 +170,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 |
@@ -403,7 +403,7 @@ to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enab
| 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` | ✅ | — | |
@@ -415,6 +415,12 @@ to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enab
| `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
diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala
index 8cdabebb85..ccca244652 100644
--- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala
@@ -644,7 +644,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.
@@ -653,14 +653,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 :+ eagerIndexReason
+
/** 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
@@ -672,7 +680,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))
}
}
}
@@ -708,7 +716,7 @@ object CometElementAt extends CometExpressionSerde[ElementAt] {
.setListExtract(arrayExtractBuilder)
.build())
} else {
- withFallbackReason(expr, "unsupported arguments for ElementAt")
+ withFallbackReason(expr, argumentsReason)
None
}
}
diff --git a/spark/src/main/scala/org/apache/comet/serde/maps.scala b/spark/src/main/scala/org/apache/comet/serde/maps.scala
index f9d3e2dbc0..a8bbdb4958 100644
--- a/spark/src/main/scala/org/apache/comet/serde/maps.scala
+++ b/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)
diff --git a/spark/src/test/resources/sql-tests/expressions/array/element_at_ansi.sql b/spark/src/test/resources/sql-tests/expressions/array/element_at_ansi.sql
index 89293ccefd..6141ebba74 100644
--- a/spark/src/test/resources/sql-tests/expressions/array/element_at_ansi.sql
+++ b/spark/src/test/resources/sql-tests/expressions/array/element_at_ansi.sql
@@ -112,14 +112,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
diff --git a/spark/src/test/resources/sql-tests/expressions/map/element_at_map.sql b/spark/src/test/resources/sql-tests/expressions/map/element_at_map.sql
index 95450b8159..acb89a423e 100644
--- a/spark/src/test/resources/sql-tests/expressions/map/element_at_map.sql
+++ b/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.
diff --git a/spark/src/test/resources/sql-tests/expressions/map/element_at_map_ansi.sql b/spark/src/test/resources/sql-tests/expressions/map/element_at_map_ansi.sql
index 1d1bececfb..66544a746b 100644
--- a/spark/src/test/resources/sql-tests/expressions/map/element_at_map_ansi.sql
+++ b/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
diff --git a/spark/src/test/resources/sql-tests/expressions/map/element_at_map_collation.sql b/spark/src/test/resources/sql-tests/expressions/map/element_at_map_collation.sql
index 763e1b31f1..d020b1bf8a 100644
--- a/spark/src/test/resources/sql-tests/expressions/map/element_at_map_collation.sql
+++ b/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
diff --git a/spark/src/test/resources/sql-tests/expressions/map/get_map_value.sql b/spark/src/test/resources/sql-tests/expressions/map/get_map_value.sql
index 2c6c23846e..7aba26ea36 100644
--- a/spark/src/test/resources/sql-tests/expressions/map/get_map_value.sql
+++ b/spark/src/test/resources/sql-tests/expressions/map/get_map_value.sql
@@ -21,7 +21,7 @@ CREATE TABLE test_map(m map<string, int>) USING parquet
statement
INSERT INTO test_map VALUES (map('a', 1, 'b', 2, 'c', 3)), (map('x', 10)), (NULL)
-query spark_answer_only
+query expect_native(getmapvalue)
SELECT m['a'], m['b'], m['c'] FROM test_map
query spark_answer_only
@@ -31,9 +31,9 @@ SELECT m['x'], m['missing'] FROM test_map
query spark_answer_only
SELECT map('a', 1, 'b', 2)['a'], map('a', 1, 'b', 2)['missing'], map('a', 1, 'b', 2)[NULL]
--- Map key types whose Spark equality Comet's native `map_extract` cannot reproduce fall back to
--- Spark. `x[key]` routes through `GetMapValue` / `CometMapExtract` (the `element_at` form in
--- `element_at_map.sql` exercises the same guard on `CometElementAt`). A `map<double,int>` column
+-- Map key types whose Spark equality Comet's native `map_extract` cannot reproduce dispatch to
+-- Spark's generated code inside Comet. `x[key]` routes through `GetMapValue` / `CometMapExtract`
+-- (`element_at_map.sql` exercises the same guard on `CometElementAt`). A `map<double,int>` column
-- is used rather than a `map(...)` literal because Spark's `SimplifyExtractValueOps` rewrites
-- `map(...)[key]` over a literal map into a `CASE` before it can reach the native lookup. Spark
-- stores a `-0.0` key as `+0.0` and finds it with `nanSafeCompareDoubles`, so it returns `7` for a
@@ -44,5 +44,5 @@ CREATE TABLE test_map_double(m map<double, int>) USING parquet
statement
INSERT INTO test_map_double VALUES (map(CAST(0 AS DOUBLE), 7)), (map(CAST(1 AS DOUBLE), 8)), (NULL)
-query expect_fallback(Spark normalizes floating-point map keys)
-SELECT m[CAST(-0.0 AS DOUBLE)] FROM test_map_double
+query expect_dispatch(getmapvalue)
+SELECT m[double('-0.0')] FROM test_map_double
diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_lookup_dispatch.sql b/spark/src/test/resources/sql-tests/expressions/map/map_lookup_dispatch.sql
new file mode 100644
index 0000000000..79477a50f7
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/map/map_lookup_dispatch.sql
@@ -0,0 +1,125 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+-- ConfigMatrix: spark.sql.ansi.enabled=false,true
+
+-- Read maps from columns so SimplifyExtractValueOps cannot rewrite m[k] into CASE.
+-- String-to-float casts preserve negative zero (a decimal -0.0 literal does not).
+statement
+CREATE TABLE lookup_double(m MAP<DOUBLE, INT>, k DOUBLE) USING parquet
+
+statement
+INSERT INTO lookup_double VALUES
+ (map(double('0.0'), 7), double('-0.0')),
+ (map(double('-0.0'), 8), double('0.0')),
+ (map(double('NaN'), 9), double('NaN')),
+ (map(double('Infinity'), 10), double('Infinity')),
+ (map(double('-Infinity'), 11), double('-Infinity')),
+ (map(double('4.9E-324'), 12), double('4.9E-324')),
+ (map(1D, 13), 2D),
+ (map(1D, NULL), 1D),
+ (map(), 1D),
+ (NULL, 1D),
+ (map(1D, 14), NULL)
+
+query expect_dispatch(getmapvalue)
+SELECT m[k] FROM lookup_double
+
+query expect_dispatch(element_at)
+SELECT element_at(m, k) FROM lookup_double
+
+-- Exercise the Float Arrow input separately, including NULL maps, keys and values.
+statement
+CREATE TABLE lookup_float(m MAP<FLOAT, INT>, k FLOAT) USING parquet
+
+statement
+INSERT INTO lookup_float VALUES
+ (map(float('0.0'), 7), float('-0.0')),
+ (map(float('NaN'), 8), float('NaN')),
+ (map(float('Infinity'), 9), float('Infinity')),
+ (map(float('1.4E-45'), 10), float('1.4E-45')),
+ (map(1F, 11), 2F), (map(1F, NULL), 1F), (map(), 1F), (NULL, 1F), (map(1F, 12), NULL)
+
+query expect_dispatch(getmapvalue)
+SELECT m[k] FROM lookup_float
+
+query expect_dispatch(element_at)
+SELECT element_at(m, k) FROM lookup_float
+
+-- Complex keys compare structurally, including nested NULLs and different lengths.
+-- Complex values exercise the dispatcher's output writer as well as key comparison.
+statement
+CREATE TABLE lookup_array(m MAP<ARRAY<INT>, ARRAY<STRING>>, k ARRAY<INT>) USING parquet
+
+statement
+INSERT INTO lookup_array VALUES
+ (map(array(1, NULL), array('a', NULL)), array(1, NULL)),
+ (map(array(1), array('b')), array(1, NULL)),
+ (map(CAST(array() AS ARRAY<INT>), CAST(array() AS ARRAY<STRING>)), CAST(array() AS ARRAY<INT>)),
+ (map(array(1), NULL), array(1)),
+ (map(), array(1)), (NULL, array(1)), (map(array(1), array('c')), NULL)
+
+query expect_dispatch(getmapvalue)
+SELECT m[k] FROM lookup_array
+
+query expect_dispatch(element_at)
+SELECT element_at(m, k) FROM lookup_array
+
+statement
+CREATE TABLE lookup_struct(m MAP<STRUCT<a: INT, b: STRING>, INT>, k STRUCT<a: INT, b: STRING>) USING parquet
+
+statement
+INSERT INTO lookup_struct VALUES
+ (map(named_struct('a', 1, 'b', NULL), 7), named_struct('a', 1, 'b', NULL)),
+ (map(named_struct('a', 1, 'b', 'a'), 8), named_struct('a', 1, 'b', 'b')),
+ (map(), named_struct('a', NULL, 'b', NULL)),
+ (NULL, named_struct('a', 1, 'b', 'a')),
+ (map(named_struct('a', 1, 'b', 'a'), 9), NULL)
+
+query expect_dispatch(getmapvalue)
+SELECT m[k] FROM lookup_struct
+
+query expect_dispatch(element_at)
+SELECT element_at(m, k) FROM lookup_struct
+
+-- Floating-point equality must also hold inside complex keys.
+statement
+CREATE TABLE lookup_nested(m MAP<ARRAY<DOUBLE>, INT>, k ARRAY<DOUBLE>) USING parquet
+
+statement
+INSERT INTO lookup_nested VALUES
+ (map(array(0D, double('NaN'), NULL), 7), array(double('-0.0'), double('NaN'), NULL)),
+ (map(array(0D), 8), array(1D))
+
+query expect_dispatch(getmapvalue)
+SELECT m[k] FROM lookup_nested
+
+query expect_dispatch(element_at)
+SELECT element_at(m, k) FROM lookup_nested
+
+-- Even with ANSI enabled, a NULL map must short-circuit a throwing key expression.
+statement
+CREATE TABLE lookup_short_circuit(id INT, m MAP<DOUBLE, INT>) USING parquet
+
+statement
+INSERT INTO lookup_short_circuit VALUES (1, map(0D, 7)), (2, NULL), (3, map(0D, 9))
+
+query expect_dispatch(getmapvalue)
+SELECT m[double(id % (id - 2))] FROM lookup_short_circuit
+
+query expect_dispatch(element_at)
+SELECT element_at(m, double(id % (id - 2))) FROM lookup_short_circuit
diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala
index a1b7a1095c..51367386b7 100644
--- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala
@@ -33,6 +33,44 @@ import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOpti
class CometMapExpressionSuite extends CometTestBase {
+ private def checkMapLookupDispatch(query: String): Unit = {
+ checkSparkAnswerAndImpl(query, native = Seq.empty, dispatched = Seq("element_at"))
+ }
+
+ test("map lookup preserves fallback when codegen dispatch is disabled") {
+ withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") {
+ withParquetTable(Seq((Map(0.0 -> 7), -0.0)), "tbl") {
+ Seq("_1[_2]", "element_at(_1, _2)").foreach { lookup =>
+ checkSparkAnswerAndFallbackReason(
+ s"SELECT $lookup FROM tbl",
+ "Spark normalizes floating-point map keys")
+ }
+ }
+ withParquetTable(Seq((Map(Seq(1) -> 7), Seq(1))), "tbl") {
+ Seq("_1[_2]", "element_at(_1, _2)").foreach { lookup =>
+ checkSparkAnswerAndFallbackReason(
+ s"SELECT $lookup FROM tbl",
+ "casts the lookup key to the map's exact Arrow key type")
+ }
+ }
+ }
+ }
+
+ test("map lookup dispatch respects expression enablement") {
+ withParquetTable(Seq((Map(0.0 -> 7), -0.0)), "tbl") {
+ Seq(
+ (CometConf.getExprEnabledConfigKey("GetMapValue"), "_1[_2]"),
+ (CometConf.getExprEnabledConfigKey("ElementAt"), "element_at(_1, _2)"))
+ .foreach { case (config, lookup) =>
+ withSQLConf(config -> "false") {
+ checkSparkAnswerAndFallbackReason(
+ s"SELECT $lookup FROM tbl",
+ "Expression support is disabled")
+ }
+ }
+ }
+ }
+
test("read map[int, int] from parquet") {
withTempPath { dir =>
@@ -389,30 +427,28 @@ class CometMapExpressionSuite extends CometTestBase {
// Finding E: a map nested inside a map value is handed to the JVM dispatcher whole, so its double
// keys never revisit `CometLiteral`. The guard has to live on the lookup instead. The outer key
- // type is INT, so inspecting only the outermost map would admit the query; the fallback comes
+ // type is INT, so inspecting only the outermost map would admit the query; dispatch comes
// from the inner `element_at`, whose child is `MapType(DoubleType, IntegerType)`. Constant folding
// is on here, which is the only configuration where the inner map becomes such a literal, so this
// regression cannot be expressed in a SQL fixture (the harness disables folding). The direct
// single-level double-key lookups live in `element_at_map.sql` / `get_map_value.sql`.
- test("nested map lookup with floating-point keys falls back") {
+ test("nested map lookup with floating-point keys dispatches") {
withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") {
val negZero = "CAST(concat('-', CAST(_1 - 1 AS STRING), '.0') AS DOUBLE)"
- checkSparkAnswerAndFallbackReason(
+ checkMapLookupDispatch(
"SELECT _1 AS id, element_at(element_at(map(1, map(CAST(0 AS DOUBLE), 7)), _1), " +
- s"$negZero) AS v FROM tbl",
- "Spark normalizes floating-point map keys")
+ s"$negZero) AS v FROM tbl")
}
}
// Finding E for a collated inner key. Same folding-on-only bypass as the double-key case above;
// the direct single-level collated lookup lives in `element_at_map_collation.sql`.
- test("nested map lookup with collated string keys falls back") {
+ test("nested map lookup with collated string keys dispatches") {
assume(isSpark40Plus)
withParquetTable(Seq(("a1", 0)), "tbl") {
- checkSparkAnswerAndFallbackReason(
+ checkMapLookupDispatch(
"SELECT element_at(element_at(map(1, map(CAST('A1' AS STRING COLLATE UTF8_LCASE), 7)), 1), " +
- "CAST(_1 AS STRING COLLATE UTF8_LCASE)) AS v FROM tbl",
- "cannot honour a non-default collation")
+ "CAST(_1 AS STRING COLLATE UTF8_LCASE)) AS v FROM tbl")
}
}
@@ -487,24 +523,22 @@ class CometMapExpressionSuite extends CometTestBase {
// literal and `element_at` with a dynamic `-0.0` lookup must match Spark's `+0.0`-normalized key.
// Native `map_extract` compares raw Arrow values, so `MapKeySupport` declines it at `element_at`.
// Spark returns 7, NULL, NULL. (`element_at_map.sql` covers the folding-off constructor form.)
- test("folded map literal with floating-point keys in element_at falls back (multirow)") {
+ test("folded map literal with floating-point keys in element_at dispatches (multirow)") {
withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") {
val lookup = "CAST(concat('-', CAST(_1 - 1 AS STRING), '.0') AS DOUBLE)"
- checkSparkAnswerAndFallbackReason(
- s"SELECT _1 AS id, element_at(map(CAST(0 AS DOUBLE), 7), $lookup) AS v FROM tbl",
- "Spark normalizes floating-point map keys")
+ checkMapLookupDispatch(
+ s"SELECT _1 AS id, element_at(map(CAST(0 AS DOUBLE), 7), $lookup) AS v FROM tbl")
}
}
// Direct single-level folded map with collated string keys. Native lookup is bytewise, so a
// case-insensitive `a1` lookup against a stored `A1` cannot match; `MapKeySupport` declines it.
- test("folded map literal with collated string keys in element_at falls back (multirow)") {
+ test("folded map literal with collated string keys in element_at dispatches (multirow)") {
assume(isSpark40Plus)
withParquetTable(Seq(("a1", 0)), "tbl") {
- checkSparkAnswerAndFallbackReason(
+ checkMapLookupDispatch(
"SELECT element_at(map(CAST('A1' AS STRING COLLATE UTF8_LCASE), 7), " +
- "CAST(_1 AS STRING COLLATE UTF8_LCASE)) AS v FROM tbl",
- "cannot honour a non-default collation")
+ "CAST(_1 AS STRING COLLATE UTF8_LCASE)) AS v FROM tbl")
}
}
@@ -512,12 +546,11 @@ class CometMapExpressionSuite extends CometTestBase {
// element; native `map_extract` casts the lookup to the key's exact Arrow type and cannot
// reproduce Spark's equality, so `MapKeySupport` declines every complex key type. Spark returns
// 7, NULL, NULL. (`element_at_map.sql` covers the constructor form with a non-null lookup.)
- test("folded map literal with complex array key in element_at falls back (multirow)") {
+ test("folded map literal with complex array key in element_at dispatches (multirow)") {
withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") {
- checkSparkAnswerAndFallbackReason(
+ checkMapLookupDispatch(
"SELECT _1 AS id, element_at(map(array(1), 7), " +
- "array(IF(_1 = 2, CAST(NULL AS INT), _1))) AS v FROM tbl",
- "casts the lookup key to the map's exact Arrow key type")
+ "array(IF(_1 = 2, CAST(NULL AS INT), _1))) AS v FROM tbl")
}
}
diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCodegenDispatchBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCodegenDispatchBenchmark.scala
index d5c2bd220e..dfb5f5009b 100644
--- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCodegenDispatchBenchmark.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCodegenDispatchBenchmark.scala
@@ -20,13 +20,21 @@
package org.apache.spark.sql.benchmark
import java.nio.charset.StandardCharsets
+import java.util.concurrent.atomic.AtomicReference
import org.apache.spark.benchmark.Benchmark
-import org.apache.spark.sql.Row
+import org.apache.spark.metrics.source.CodegenMetrics
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference, ElementAt, Expression, GetMapValue}
+import org.apache.spark.sql.catalyst.expressions.codegen.CodeGenerator
import org.apache.spark.sql.catalyst.optimizer.ConstantFolding
+import org.apache.spark.sql.comet.{CometPlan, CometProjectExec}
+import org.apache.spark.sql.execution.{CommandResultExec, ProjectExec, QueryExecution, SparkPlan}
+import org.apache.spark.sql.execution.datasources.v2.V2TableWriteExec
import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.util.QueryExecutionListener
-import org.apache.comet.CometConf
+import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo}
import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus
import org.apache.comet.udf.codegen.CometScalaUDFCodegen
@@ -52,19 +60,28 @@ import org.apache.comet.udf.codegen.CometScalaUDFCodegen
* same work as the first row, so the spread between the two is this machine's noise floor for
* that table; ignore any difference between the other rows that is smaller than that spread.
*
- * Compilation is a one-time cost, and the steady-state tables warm up before timing, so it does
- * not appear there. The `first use` table at the top measures it directly: the same query run
- * twice back to back on a cold kernel, after the dispatcher itself has been warmed on an
- * unrelated expression.
+ * Steady-state tables warm JVM code/source/JIT state before timing but still include per-task
+ * kernel setup. First/second execution times also include class loading, I/O and scheduling;
+ * their difference is not an isolated compilation cost. The general cases can share compiled
+ * source, while the map first-use mode below isolates each selected case in a fresh JVM.
*
- * Before timing anything, every case is run through all three arms and the rows are compared, so
- * a timing cannot come from an arm that computed something else.
+ * Before warmed timing, every case is run through all three arms and the rows are compared.
+ * First-use measurements precede validation to avoid populating the source cache; discard their
+ * results if the subsequent checks fail.
*
* To run this benchmark:
* {{{
* SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometCodegenDispatchBenchmark
* }}}
* Results will be written to "spark/benchmarks/CometCodegenDispatchBenchmark-**results.txt".
+ *
+ * Map lookups have a separate, strict-checked matrix. Pass `--map-lookups` for warmed tables,
+ * optionally `--case=get_map_value-double-4-lookup` to select one case or `--check-only` to
+ * validate without timing. For first use, launch a fresh JVM per case and arm with, for example,
+ * `--map-lookups --case=get_map_value-double-4-lookup --first-use=dispatch`. The other arms are
+ * `fallback` and `spark`. This measures first and second query execution on 1024 rows before
+ * answer/route checks; it does not run the warmed tables. Never combine cold cases in one JVM:
+ * map sizes and native sibling expressions can share the same compiled lookup source.
*/
object CometCodegenDispatchBenchmark extends CometBenchmarkBase {
@@ -167,6 +184,11 @@ object CometCodegenDispatchBenchmark extends CometBenchmarkBase {
private val SparkCaseName = "Spark (Comet disabled)"
override def runCometBenchmark(mainArgs: Array[String]): Unit = {
+ if (mainArgs.headOption.contains("--map-lookups")) {
+ runMapLookups(mainArgs.drop(1))
+ return
+ }
+ require(mainArgs.isEmpty, s"Unknown arguments: ${mainArgs.mkString(" ")}")
val selected = cases
runBenchmark("Codegen dispatch: environment") {
emitEnvironment(selected)
@@ -175,9 +197,8 @@ object CometCodegenDispatchBenchmark extends CometBenchmarkBase {
return
}
- // The kernels must be uncompiled when `runFirstUse` starts, so it runs before anything else
- // executes these queries: `CodeGenerator.compile` caches on the generated source JVM-wide,
- // and a single earlier execution would make every "first use" number a cache hit.
+ // Run before answer checks populate the JVM-wide source cache. Some general cases still
+ // share sources with each other; the map-specific first-use mode avoids that by selection.
withCorpus(SmallRows) {
runBenchmark(s"Codegen dispatch: first use, $SmallRows rows") {
runFirstUse(selected)
@@ -189,13 +210,252 @@ object CometCodegenDispatchBenchmark extends CometBenchmarkBase {
selected.foreach(verifyArmsAgree(_, LargeRows))
selected.foreach(runSteadyState(_, LargeRows))
}
+ runMapLookups(Array.empty)
+ }
+
+ // Bound the nested-column corpus size while retaining multiple Comet batches per task.
+ private val MapLargeRows = 64 * 1024
+
+ private case class MapCase(lookup: String, keyType: String, entries: Int, mixed: Boolean) {
+ val name: String = s"$lookup-$keyType-$entries-${if (mixed) "mixed" else "lookup"}"
+ val c: DispatchCase = {
+ val value = if (lookup == "get_map_value") "m[k]" else "element_at(m, k)"
+ val siblings = if (mixed) ", length(c_str), c_long + 1, substring(c_str, 1, 4)" else ""
+ DispatchCase(
+ name,
+ s"SELECT $value AS v$siblings FROM parquetV1Table",
+ extraConfigs = Seq(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false"))
+ }
+
+ def isLookup(e: Expression): Boolean = e match {
+ case _: GetMapValue => lookup == "get_map_value"
+ case _: ElementAt => lookup == "element_at"
+ case _ => false
+ }
+ }
+
+ private def mapCases: Seq[MapCase] = for {
+ keyType <- Seq("double", "struct")
+ entries <- Seq(4, 64)
+ lookup <- Seq("get_map_value", "element_at")
+ mixed <- Seq(false, true)
+ } yield MapCase(lookup, keyType, entries, mixed)
+
+ private def mapArms(c: DispatchCase): Seq[(String, Seq[(String, String)])] = Seq(
+ "dispatch" -> dispatchConfigs(c),
+ "fallback" -> fallbackConfigs(c),
+ "spark" -> sparkConfigs(c))
+
+ private def runMapLookups(args: Array[String]): Unit = {
+ require(
+ args.forall(a =>
+ a.startsWith("--case=") || a.startsWith("--first-use=") ||
+ a == "--check-only"),
+ s"Unknown map arguments: ${args.mkString(" ")}")
+ def option(prefix: String): Option[String] = {
+ val values = args.filter(_.startsWith(prefix))
+ require(values.length <= 1, s"Repeated option: $prefix")
+ values.headOption.map(_.stripPrefix(prefix))
+ }
+ val filter = option("--case=")
+ val firstUse = option("--first-use=")
+ val checkOnly = args.contains("--check-only")
+ val selected = mapCases.filter(c => filter.forall(_ == c.name))
+ require(
+ selected.nonEmpty,
+ s"Unknown case; choose from ${mapCases.map(_.name).mkString(", ")}")
+ require(
+ firstUse.isEmpty || (selected.size == 1 && !checkOnly),
+ "First use requires exactly one --case and no --check-only; launch a fresh JVM per arm.")
+ require(firstUse.forall(Set("dispatch", "fallback", "spark")), "Unknown first-use arm")
+ require(
+ spark.sparkContext.isLocal,
+ "Map route counters and codegen metrics require local mode")
+
+ runBenchmark("Map lookup dispatch: environment") {
+ emitEnvironment(selected.map(_.c), MapLargeRows)
+ emit("First-use mode measures only the sub-batch corpus; it does not run warmed tables.")
+ emit("Keys: DOUBLE or STRUCT<k:INT,tag:INT>; values: BIGINT; 4 or 64 entries.")
+ emit("Column lookups: 25% last-entry hit, 25% miss, 25% NULL value, 25% first-entry hit.")
+ emit("DOUBLE first-entry lookups use -0.0 against +0.0. Map construction is not timed.")
+ emit("Timings include planning, Parquet scan and noop sink; no Spark data cache is used.")
+ emit("Warmed queries still create tasks/kernels; only JVM code/source/JIT state is warmed.")
+ }
+
+ firstUse match {
+ case Some(arm) =>
+ val c = selected.head
+ withMapCorpus(c, SmallRows) {
+ runMapFirstUse(c, arm)
+ verifyMapCase(c, SmallRows)
+ }
+ case None =>
+ for (rows <- Seq(SmallRows, MapLargeRows);
+ keyType <- Seq("double", "struct"); entries <- Seq(4, 64)) {
+ val group = selected.filter(c => c.keyType == keyType && c.entries == entries)
+ group.headOption.foreach { first =>
+ withMapCorpus(first, rows) {
+ group.foreach { c =>
+ verifyMapCase(c, rows)
+ if (!checkOnly) runSteadyState(c.c, rows)
+ }
+ }
+ }
+ }
+ }
+ }
+
+ private def withMapCorpus(c: MapCase, rows: Int)(f: => Unit): Unit = {
+ def key(ordinal: String): String = if (c.keyType == "double") {
+ s"CAST($ordinal AS DOUBLE)"
+ } else {
+ s"named_struct('k', CAST($ordinal AS INT), 'tag', CAST(pmod(id, 7) AS INT))"
+ }
+ val lookupKey = if (c.keyType == "double") {
+ "IF(j = 0, CAST('-0.0' AS DOUBLE), CAST(j AS DOUBLE))"
+ } else key("j")
+ val query = s"""
+ |SELECT map_from_arrays(
+ | transform(sequence(0, ${c.entries - 1}), x -> ${key("x")}),
+ | transform(sequence(0, ${c.entries - 1}),
+ | x -> IF(x = 1, CAST(NULL AS BIGINT), id * 128 + x))) AS m,
+ | $lookupKey AS k, id AS c_long, repeat(CAST(id AS STRING), 4) AS c_str
+ |FROM (SELECT id, CASE pmod(id, 4)
+ | WHEN 0 THEN ${c.entries - 1} WHEN 1 THEN ${c.entries}
+ | WHEN 2 THEN 1 ELSE 0 END AS j FROM $tbl)
+ |""".stripMargin
+ withTempPath { dir =>
+ withTempTable(tbl, "parquetV1Table") {
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark.range(rows).createOrReplaceTempView(tbl)
+ prepareTable(dir, spark.sql(query))
+ }
+ f
+ }
+ }
+ }
+
+ private def runMapFirstUse(c: MapCase, arm: String): Unit = {
+ // Warm common machinery without executing any lookup. Neither answer checks nor another
+ // map case may precede the selected first use: CodeGenerator's source cache is JVM-wide.
+ val warmup = DispatchCase("warmup", "SELECT c_str rlike '[0-9]+' FROM parquetV1Table")
+ mapArms(warmup).foreach { case (_, configs) =>
+ (0 until 5).foreach(_ => runQuery(warmup.query, configs))
+ }
+ val configs = mapArms(c.c).find(_._1 == arm).get._2
+ runBenchmark(s"${c.name}: first use ($arm), $SmallRows rows") {
+ (1 to 2).foreach { iteration =>
+ CometScalaUDFCodegen.resetStats()
+ val count = CodegenMetrics.METRIC_COMPILATION_TIME.getCount
+ val nanos = CodeGenerator.compileTime
+ val elapsed = timeMillis(runQuery(c.c.query, configs))
+ val compilationMs = (CodeGenerator.compileTime - nanos) / 1e6
+ val compilations = CodegenMetrics.METRIC_COMPILATION_TIME.getCount - count
+ val stats = CometScalaUDFCodegen.stats()
+ emit(f"FIRST_USE case=${c.name} arm=$arm run=$iteration wall_ms=$elapsed%.3f " +
+ f"jvm_compilations=$compilations jvm_compile_ms=$compilationMs%.3f " +
+ s"task_kernel_initializations=${stats.compileCount} batch_cache_hits=${stats.cacheHitCount}")
+ }
+ emit("JVM compilation metrics include all Spark-generated classes, not only Comet kernels.")
+ emit("First-minus-second is not compilation time: class loading, JIT and I/O also differ.")
+ }
+ }
+
+ private def verifyMapCase(c: MapCase, rows: Int): Unit = {
+ verifyArmsAgree(c.c, rows)
+ val explain = new ExtendedExplainInfo
+ mapArms(c.c).foreach { case (arm, configs) =>
+ withSQLConf(configs: _*) {
+ val df = spark.sql(c.c.query)
+ val lookups = df.queryExecution.optimizedPlan.flatMap(_.expressions.flatMap(_.collect {
+ case e if c.isLookup(e) => e
+ }))
+ assert(lookups.size == 1, s"${c.name}: lookup was optimized away or duplicated")
+ assert(
+ lookups.head.children.forall(_.isInstanceOf[AttributeReference]),
+ s"${c.name}: lookup arguments must remain persisted columns")
+ val name = CometExplainInfo.exprDisplayName(lookups.head)
+ CometScalaUDFCodegen.resetStats()
+ val plan = noopInputPlan(df)
+ val stats = CometScalaUDFCodegen.stats()
+ val dispatched = explain.getCodegenDispatchExpressions(plan).toSet
+ val native = explain.getNativeExpressions(plan).toSet
+ val activity = stats.compileCount + stats.cacheHitCount
+ if (arm == "dispatch") {
+ assert(plan.exists(_.isInstanceOf[CometProjectExec]), s"Missing CometProject: $plan")
+ assert(!plan.exists(_.isInstanceOf[ProjectExec]), s"Spark projection remained: $plan")
+ assert(findFirstNonCometOperator(plan).isEmpty, s"Unexpected operator fallback: $plan")
+ assert(
+ dispatched == Set(name) && !native.contains(name),
+ s"${c.name}: dispatched=$dispatched, native=$native")
+ assert(activity > 0, s"${c.name}: dispatcher did not execute")
+ if (c.mixed) {
+ assert(
+ Set("length", "add", "substring").subsetOf(native),
+ s"${c.name}: mixed siblings must be native, found $native")
+ }
+ } else {
+ assert(
+ plan.exists {
+ case p: ProjectExec => p.projectList.exists(_.exists(c.isLookup))
+ case _ => false
+ },
+ s"${c.name}: missing Spark lookup projection: $plan")
+ assert(activity == 0 && dispatched.isEmpty, s"$arm unexpectedly dispatched")
+ if (arm == "spark") {
+ assert(!plan.exists(_.isInstanceOf[CometPlan]), s"Pure Spark contained Comet: $plan")
+ } else {
+ assert(plan.exists(_.isInstanceOf[CometPlan]), s"Fallback arm lost Comet scan: $plan")
+ }
+ }
+ // This separate collect also verifies that the corpus is consumed in its entirety.
+ val answers = df.collect()
+ assert(
+ answers.length == rows && answers.count(!_.isNullAt(0)) == rows / 2,
+ s"${c.name}: unexpected corpus cardinality or hit/NULL distribution")
+ emit(
+ s"PASS map answers/routes: ${c.name}, rows=$rows, arm=$arm, " +
+ s"dispatched=${dispatched.toSeq.sorted.mkString(",")}, " +
+ s"native=${native.toSeq.sorted.mkString(",")}, " +
+ s"task_kernel_initializations=${stats.compileCount}, batch_cache_hits=${stats.cacheHitCount}")
+ }
+ }
+ }
+
+ /** Validate the actual noop write's input, which may differ from a standalone SELECT plan. */
+ private def noopInputPlan(df: DataFrame): SparkPlan = {
+ val captured = new AtomicReference[SparkPlan]()
+ val listener = new QueryExecutionListener {
+ override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = {
+ captured.set(qe.executedPlan)
+ }
+ override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit =
+ ()
+ }
+ // Listener synchronization is for validation only, never part of a timed benchmark case.
+ spark.sparkContext.listenerBus.waitUntilEmpty(10000)
+ spark.listenerManager.register(listener)
+ try {
+ df.noop()
+ spark.sparkContext.listenerBus.waitUntilEmpty(10000)
+ } finally {
+ spark.listenerManager.unregister(listener)
+ }
+ assert(captured.get() != null, "No executed noop query was captured")
+ val executed = captured.get() match {
+ case c: CommandResultExec => c.commandPhysicalPlan
+ case other => other
+ }
+ val inputs = collect(executed) { case w: V2TableWriteExec => w.query }
+ assert(inputs.size == 1, s"Expected one noop write input, found: $executed")
+ stripAQEPlan(inputs.head)
}
/**
* The reader of a results file cannot see the confs the numbers were produced under, and for
* this benchmark the batch size and the dispatcher conf are the whole point.
*/
- private def emitEnvironment(selected: Seq[DispatchCase]): Unit = {
+ private def emitEnvironment(selected: Seq[DispatchCase], largeRows: Int = LargeRows): Unit = {
emit(s"Spark version: ${spark.version}")
emit(
s"Java version: ${System.getProperty("java.version")} " +
@@ -213,7 +473,7 @@ object CometCodegenDispatchBenchmark extends CometBenchmarkBase {
emit(
" at least 2 iterations and at least 2s of timed iterations; the table reports best and")
emit(" average of the timed iterations.")
- emit(s"Row counts: $SmallRows (sub-batch) and $LargeRows (multi-batch).")
+ emit(s"Row counts: $SmallRows (sub-batch) and $largeRows (multi-batch).")
val skipped =
Seq("to_time(fmt)" -> isSpark41Plus).collect { case (name, false) => name }
if (skipped.nonEmpty) {
@@ -225,15 +485,9 @@ object CometCodegenDispatchBenchmark extends CometBenchmarkBase {
}
/**
- * Cost of the first execution of a query whose kernel has never been compiled, against the
- * second execution of the same query.
- *
- * The dispatcher is warmed on an expression that is not in the case list first, so that Janino,
- * `CodeGenerator`, the Arrow bridge and the FFI boundary are all loaded and JIT-warm before the
- * first case runs and the difference is dominated by compiling that one kernel. It is still an
- * upper bound on the compile: the first case in the list absorbs whatever class loading the
- * warmup missed, and any case that reaches machinery no earlier case did -- the grouped case is
- * the first to shuffle -- pays for that here too.
+ * Observed first and second executions after unrelated dispatcher warmup. General cases may
+ * share compiled sources, and the timings include work other than compilation. Use the
+ * selected-case map mode for isolated first-use measurements and actual compilation metrics.
*/
private def runFirstUse(selected: Seq[DispatchCase]): Unit = {
// `rlike` routes through the same dispatcher and is not one of the cases below.
@@ -242,7 +496,7 @@ object CometCodegenDispatchBenchmark extends CometBenchmarkBase {
emit(
f"${"case"}%-24s ${"1st run (ms)"}%14s ${"2nd run (ms)"}%14s " +
- f"${"one-time (ms)"}%14s ${"compiles"}%9s")
+ f"${"1st - 2nd (ms)"}%14s ${"compiles"}%9s")
emit("-" * 84)
selected.foreach { c =>
CometScalaUDFCodegen.resetStats()
@@ -256,7 +510,8 @@ object CometCodegenDispatchBenchmark extends CometBenchmarkBase {
emit("itself is deduplicated JVM-wide by Spark's CodeGenerator source cache, so only the")
emit("first task to reach a given kernel source pays for it -- and a case whose bound tree")
emit("matches an earlier case's, as `to_binary(utf-8)` matches `encode(utf-8)`, is a hit on")
- emit("that cache and reports a one-time cost near zero.")
+ emit(
+ "that cache. First-minus-second can also reflect JIT, I/O, scheduling and class loading.")
}
/**
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org