From 44411f1e61fd5a8dbe95e9c1093527f31633121d Mon Sep 17 00:00:00 2001 From: LinSimon-901101 Date: Fri, 11 Sep 2026 00:17:40 +0800 Subject: [PATCH 1/2] feat: route map_from_arrays LAST_WIN through codegen dispatcher --- .../expression-audits/map_funcs.md | 10 ++++++---- docs/source/user-guide/latest/expressions.md | 2 +- .../scala/org/apache/comet/serde/maps.scala | 4 +++- .../map/map_from_arrays_dedup_policy.sql | 18 ++++++++++-------- 4 files changed, 20 insertions(+), 14 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/map_funcs.md b/docs/source/contributor-guide/expression-audits/map_funcs.md index ea13e6ab130..b340ee24f4c 100644 --- a/docs/source/contributor-guide/expression-audits/map_funcs.md +++ b/docs/source/contributor-guide/expression-audits/map_funcs.md @@ -44,10 +44,12 @@ ## map_from_arrays -- Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. -- Spark 3.5.8 (audited 2026-05-27): baseline. `MapFromArrays(left, right) extends BinaryExpression with NullIntolerant`; Spark uses `ArrayBasedMapBuilder` to detect duplicate keys (subject to `spark.sql.mapKeyDedupPolicy`) and rejects null keys with `RuntimeException("Cannot use null as map key")`. Comet `CometMapFromArrays` wraps the inputs in `CaseWhen(IsNotNull(left) AND IsNotNull(right), map(left, right), null)` so NULL-array inputs return NULL rather than triggering the previously reported native crash ([#3327](https://github.com/apache/datafusion-comet/issues/3327)). -- Spark 4.0.1 (audited 2026-05-27): semantics unchanged; `NullIntolerant` trait replaced by `nullIntolerant: Boolean`. -- Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. +- Spark 3.4.3 (audited 2026-09-11): baseline. `MapFromArrays(left, right)` accepts two arrays and returns a `MapType` whose key and value types come from the corresponding array element types. Both interpreted evaluation and generated code copy the arrays and call `ArrayBasedMapBuilder.from`. A NULL input array returns NULL. A NULL key element raises `NULL_MAP_KEY`, unequal array lengths raise an error, duplicate keys raise `DUPLICATED_MAP_KEY` under the default `EXCEPTION` policy, and `LAST_WIN` keeps the last value at the key's first insertion position. +- Spark 3.5.8 (audited 2026-09-11): runtime behavior and generated code are unchanged. Adds `stateful = true` because the expression reuses a mutable `ArrayBasedMapBuilder`. +- Spark 4.0.1 (audited 2026-09-11): replaces the `NullIntolerant` trait with `nullIntolerant = true`. `ArrayBasedMapBuilder` adds floating-point key normalization and collation-aware string-key equality. `spark.sql.legacy.disableMapKeyNormalization=true` restores distinct `-0.0` and `+0.0` keys. `VariantType` is also rejected as a map key. The expression's interpreted and generated paths remain unchanged. +- Spark 4.1.1 (audited 2026-09-11): adds the `MAP_FROM_ARRAYS` tree pattern and changes `spark.sql.mapKeyDedupPolicy` from an internal string config to an enum config. Runtime behavior and generated code are unchanged. +- Comet runs the default `EXCEPTION` configuration natively and guards NULL-array inputs before native map construction ([#3327](https://github.com/apache/datafusion-comet/issues/3327)). Under `LAST_WIN`, `CodegenDispatchFallback` routes `MapFromArrays` through Spark's generated code inside the Comet pipeline by default. Setting `spark.comet.expression.MapFromArrays.allowIncompatible=true` opts into the divergent native path; if the dispatcher is disabled instead, the enclosing operator falls back to Spark. +- Known limitation: the native path still does not reject a NULL element inside the keys array ([#4680](https://github.com/apache/datafusion-comet/issues/4680)). The `LAST_WIN` dispatcher path happens to reject it through Spark's builder, but this does not resolve the default native-path divergence. ## map_from_entries diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 231d80affbf..da1fde4dabe 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -400,7 +400,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci | `map_concat` | ✅ | Codegen dispatch | | | `map_contains_key` | ✅ | — | | | `map_entries` | ✅ | Native | | -| `map_from_arrays` | ✅ | Native | | +| `map_from_arrays` | ✅ | Hybrid | LAST_WIN routes through the JVM codegen dispatcher; EXCEPTION runs natively | | `map_from_entries` | ✅ | Hybrid | BinaryType key/value falls back (Incompatible) ([details](compatibility/expressions/map.md)) | | `map_keys` | ✅ | Native | | | `map_values` | ✅ | Native | | 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 51fa428b543..1443f92ec54 100644 --- a/spark/src/main/scala/org/apache/comet/serde/maps.scala +++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala @@ -151,7 +151,9 @@ private object MapKeyDedupPolicySupport { .equalsIgnoreCase(SQLConf.MapKeyDedupPolicy.LAST_WIN.toString) } -object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { +object CometMapFromArrays + extends CometExpressionSerde[MapFromArrays] + with CodegenDispatchFallback { override def getIncompatibleReasons(): Seq[String] = Seq(MapKeyDedupPolicySupport.incompatibleReason) diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_dedup_policy.sql b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_dedup_policy.sql index fffaf5f9a92..4cc633bf00b 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_dedup_policy.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_dedup_policy.sql @@ -15,12 +15,14 @@ -- specific language governing permissions and limitations -- under the License. --- Verifies that `map_from_arrays` falls back to Spark when `spark.sql.mapKeyDedupPolicy` is set --- to `LAST_WIN`. Spark's ArrayBasedMapBuilder keeps the last occurrence of each duplicate key; --- Comet's native `map` scalar has no LAST_WIN path, so it must fall back. The default `EXCEPTION` --- mode agrees with Comet and is covered by `map_from_arrays.sql`. +-- Verifies that `map_from_arrays` routes through the JVM codegen dispatcher when +-- `spark.sql.mapKeyDedupPolicy` is set to `LAST_WIN`. Spark's ArrayBasedMapBuilder keeps the last +-- occurrence of each duplicate key; Comet's native `map` scalar has no LAST_WIN path, so the +-- dispatcher runs Spark's generated code inside the Comet pipeline. The default `EXCEPTION` mode +-- agrees with Comet's native implementation and is covered by `map_from_arrays.sql`. -- Config: spark.sql.mapKeyDedupPolicy=LAST_WIN +-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true statement CREATE TABLE test_map_from_arrays_dedup(k array, v array) USING parquet @@ -31,11 +33,11 @@ INSERT INTO test_map_from_arrays_dedup VALUES (array('a', 'a', 'b'), array(1, 2, 3)), (array('x', 'x'), array(10, 20)) --- literal duplicate keys under LAST_WIN: Spark keeps the last value; Comet must fall back. -query expect_fallback(mapKeyDedupPolicy) +-- literal duplicate keys under LAST_WIN: Spark's generated code keeps the last value. +query expect_dispatch(map_from_arrays) SELECT map_from_arrays(array('a', 'a', 'b'), array(1, 2, 3)) --- column input falls back the same way; the incompat branch is triggered by the SQLConf value, +-- column input dispatches the same way; the incompat branch is triggered by the SQLConf value, -- not per-row content. -query expect_fallback(mapKeyDedupPolicy) +query expect_dispatch(map_from_arrays) SELECT map_from_arrays(k, v) FROM test_map_from_arrays_dedup From 16c4afef59d9d35dc533d38f497cee23257d1f3e Mon Sep 17 00:00:00 2001 From: LinSimon-901101 Date: Fri, 11 Sep 2026 22:37:19 +0800 Subject: [PATCH 2/2] test: expand map_from_arrays dispatch coverage and add benchmark --- .../expression-audits/map_funcs.md | 4 +- .../map/map_from_arrays_dedup_policy.sql | 44 +++- .../comet/CometMapExpressionSuite.scala | 37 ++- .../CometMapFromArraysBenchmark.scala | 238 ++++++++++++++++++ 4 files changed, 320 insertions(+), 3 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapFromArraysBenchmark.scala diff --git a/docs/source/contributor-guide/expression-audits/map_funcs.md b/docs/source/contributor-guide/expression-audits/map_funcs.md index b340ee24f4c..58e4b9ef18a 100644 --- a/docs/source/contributor-guide/expression-audits/map_funcs.md +++ b/docs/source/contributor-guide/expression-audits/map_funcs.md @@ -45,10 +45,12 @@ ## map_from_arrays - Spark 3.4.3 (audited 2026-09-11): baseline. `MapFromArrays(left, right)` accepts two arrays and returns a `MapType` whose key and value types come from the corresponding array element types. Both interpreted evaluation and generated code copy the arrays and call `ArrayBasedMapBuilder.from`. A NULL input array returns NULL. A NULL key element raises `NULL_MAP_KEY`, unequal array lengths raise an error, duplicate keys raise `DUPLICATED_MAP_KEY` under the default `EXCEPTION` policy, and `LAST_WIN` keeps the last value at the key's first insertion position. -- Spark 3.5.8 (audited 2026-09-11): runtime behavior and generated code are unchanged. Adds `stateful = true` because the expression reuses a mutable `ArrayBasedMapBuilder`. +- Spark 3.5.8 (audited 2026-09-11): the expression's evaluation logic and generated code are unchanged. Adds `stateful = true` because the expression reuses a mutable `ArrayBasedMapBuilder`. The underlying columnar array copy differs from 3.4.3 for NULL values; see the reader caveat below. - Spark 4.0.1 (audited 2026-09-11): replaces the `NullIntolerant` trait with `nullIntolerant = true`. `ArrayBasedMapBuilder` adds floating-point key normalization and collation-aware string-key equality. `spark.sql.legacy.disableMapKeyNormalization=true` restores distinct `-0.0` and `+0.0` keys. `VariantType` is also rejected as a map key. The expression's interpreted and generated paths remain unchanged. - Spark 4.1.1 (audited 2026-09-11): adds the `MAP_FROM_ARRAYS` tree pattern and changes `spark.sql.mapKeyDedupPolicy` from an internal string config to an enum config. Runtime behavior and generated code are unchanged. - Comet runs the default `EXCEPTION` configuration natively and guards NULL-array inputs before native map construction ([#3327](https://github.com/apache/datafusion-comet/issues/3327)). Under `LAST_WIN`, `CodegenDispatchFallback` routes `MapFromArrays` through Spark's generated code inside the Comet pipeline by default. Setting `spark.comet.expression.MapFromArrays.allowIncompatible=true` opts into the divergent native path; if the dispatcher is disabled instead, the enclosing operator falls back to Spark. +- The `LAST_WIN` SQL fixture pins dispatcher execution for literal and column inputs, including NULL arrays, empty arrays, NULL values, and a duplicate key whose last value is NULL. It also checks NULL-key and unequal-length errors alongside a positive dispatcher sentinel. `CometMapExpressionSuite` additionally checks each error query's dispatch tag, runtime dispatcher activity, and Spark exception type, error class, and SQLSTATE parity. +- Spark 3.4.3's vectorized reader can lose primitive-array NULLs in `ColumnarArray.copy` ([SPARK-48019](https://issues.apache.org/jira/browse/SPARK-48019)). The `LAST_WIN` SQL fixture disables Spark's nested-column vectorized reader so its reference preserves NULL values, while still requiring Comet operators and dispatcher execution. This is a test-baseline workaround, not a change to Comet's production reader settings. - Known limitation: the native path still does not reject a NULL element inside the keys array ([#4680](https://github.com/apache/datafusion-comet/issues/4680)). The `LAST_WIN` dispatcher path happens to reject it through Spark's builder, but this does not resolve the default native-path divergence. ## map_from_entries diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_dedup_policy.sql b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_dedup_policy.sql index 4cc633bf00b..7e073b355e4 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_dedup_policy.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_dedup_policy.sql @@ -23,6 +23,10 @@ -- Config: spark.sql.mapKeyDedupPolicy=LAST_WIN -- Config: spark.comet.exec.scalaUDF.codegen.enabled=true +-- Spark 3.4.3's ColumnarArray.copy drops primitive-array NULLs (SPARK-48019). +-- Use Spark's row reader as the reference; Comet still scans and dispatches this expression. +-- https://issues.apache.org/jira/browse/SPARK-48019 +-- Config: spark.sql.parquet.enableNestedColumnVectorizedReader=false statement CREATE TABLE test_map_from_arrays_dedup(k array, v array) USING parquet @@ -31,7 +35,15 @@ statement INSERT INTO test_map_from_arrays_dedup VALUES (array('a', 'b', 'c'), array(1, 2, 3)), (array('a', 'a', 'b'), array(1, 2, 3)), - (array('x', 'x'), array(10, 20)) + (array('x', 'x'), array(10, 20)), + -- NULL arrays stay as column inputs so NullPropagation cannot fold away map_from_arrays. + (NULL, array(1)), + (array('a'), NULL), + (NULL, NULL), + (cast(array() as array), cast(array() as array)), + -- NULL values are allowed, including the last value of a duplicate key. + (array('a', 'b'), array(NULL, 2)), + (array('a', 'a', 'b'), array(1, NULL, 3)) -- literal duplicate keys under LAST_WIN: Spark's generated code keeps the last value. query expect_dispatch(map_from_arrays) @@ -41,3 +53,33 @@ SELECT map_from_arrays(array('a', 'a', 'b'), array(1, 2, 3)) -- not per-row content. query expect_dispatch(map_from_arrays) SELECT map_from_arrays(k, v) FROM test_map_from_arrays_dedup + +-- Typed empty literal arrays exercise the dispatcher without introducing NullType inputs. +query expect_dispatch(map_from_arrays) +SELECT map_from_arrays(cast(array() as array), cast(array() as array)) + +statement +CREATE TABLE test_map_from_arrays_errors(id int, k array, v array) USING parquet + +statement +INSERT INTO test_map_from_arrays_errors VALUES + (0, array('a', 'a'), array(1, 2)), + (1, array('a', NULL), array(1, 2)), + (2, array('a', 'b'), array(1)), + (3, array('a'), array(1, 2)) + +-- Positive sentinel for the error queries: identical column types and LAST_WIN configuration +-- must dispatch. CometMapExpressionSuite also pins dispatch and exception parity for each error. +query expect_dispatch(map_from_arrays) +SELECT map_from_arrays(k, v) FROM test_map_from_arrays_errors WHERE id = 0 + +-- LAST_WIN does not permit NULL keys; this tests Spark's builder, not the native #4680 path. +query expect_error(NULL_MAP_KEY) +SELECT map_from_arrays(k, v) FROM test_map_from_arrays_errors WHERE id = 1 + +-- Both directions of the array-length mismatch must fail. +query expect_error(The key array and value array of MapData must have the same length) +SELECT map_from_arrays(k, v) FROM test_map_from_arrays_errors WHERE id = 2 + +query expect_error(The key array and value array of MapData must have the same length) +SELECT map_from_arrays(k, v) FROM test_map_from_arrays_errors WHERE id = 3 diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index f4a559b872b..5816b9608a6 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -31,7 +31,7 @@ import org.apache.spark.sql.types.BinaryType import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions} -class CometMapExpressionSuite extends CometTestBase { +class CometMapExpressionSuite extends CometTestBase with CometCodegenAssertions { test("read map[int, int] from parquet") { @@ -126,6 +126,41 @@ class CometMapExpressionSuite extends CometTestBase { } } + test("map_from_arrays LAST_WIN dispatcher preserves builder errors") { + withSQLConf( + SQLConf.MAP_KEY_DEDUP_POLICY.key -> "LAST_WIN", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + withTable("test_map_from_arrays_errors") { + sql( + "CREATE TABLE test_map_from_arrays_errors " + + "(id INT, k ARRAY, v ARRAY) USING parquet") + sql( + "INSERT INTO test_map_from_arrays_errors VALUES " + + "(1, array('a', NULL), array(1, 2)), " + + "(2, array('a', 'b'), array(1)), " + + "(3, array('a'), array(1, 2))") + + Seq(1 -> "NULL_MAP_KEY", 2 -> "_LEGACY_ERROR_TEMP_2128", 3 -> "_LEGACY_ERROR_TEMP_2128") + .foreach { case (id, errorClass) => + withClue(s"map_from_arrays error case $id: ") { + val df = sql( + "SELECT map_from_arrays(k, v) " + + s"FROM test_map_from_arrays_errors WHERE id = $id") + // The SQL fixture checks error messages, but an error alone could pass on Spark + // fallback. Pin the error query's route and actual dispatcher activity as well. + assertExpressionImpl( + df.queryExecution.executedPlan, + native = Seq.empty, + dispatched = Seq("map_from_arrays")) + assertCodegenRan { + checkSparkError(df, errorClass) + } + } + } + } + } + } + test("size with map input") { withTempDir { dir => withTempView("t1") { diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapFromArraysBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapFromArraysBenchmark.scala new file mode 100644 index 00000000000..7d67766621d --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapFromArraysBenchmark.scala @@ -0,0 +1,238 @@ +/* + * 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. + */ + +package org.apache.spark.sql.benchmark + +import java.nio.charset.StandardCharsets + +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.expressions.MapFromArrays +import org.apache.spark.sql.catalyst.optimizer.ConstantFolding +import org.apache.spark.sql.execution.{ProjectExec, SparkPlan} +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.{CometConf, ExtendedExplainInfo} +import org.apache.comet.udf.codegen.CometScalaUDFCodegen + +/** + * Measures retaining the Comet projection via the LAST_WIN codegen dispatcher against falling + * that projection back to Spark. Pure Spark is a third reference. The incompatible native + * implementation is disabled in every arm. Unique/duplicate keys, arrays of length 4/256, and + * standalone/mixed projections form eight tables. + * + * Inputs are varying String-key/Long-value columns in prepared Parquet, not constant arrays. The + * optional argument is the input-element budget per table (default 1048576): large arrays use + * fewer rows to bound memory. Compare arms within a table, not row rates across array sizes. Data + * generation, full-result comparisons and plan/runtime checks are outside timing. Timing uses + * noop(), which materializes the projection, and includes scanning and query execution. + * + * Spark's Benchmark warms each arm for 2 seconds, then times at least 2 iterations and at least 2 + * seconds. This measures steady state, not first-use compilation. The repeated dispatch-off + * baseline at the end helps expose drift; differences smaller than that spread or the reported + * standard deviation are not evidence of a speedup. Repeat the suite for independent samples. + * + * Run from the repository root (the make target builds a release library first): + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometMapFromArraysBenchmark BENCH_HEAP=4g + * }}} + * Append `-- 262144` for a smaller corpus. Results are written under spark/benchmarks/. The + * shared benchmark base uses local[1], so runtime dispatcher counters are visible to the driver. + */ +object CometMapFromArraysBenchmark extends CometBenchmarkBase { + + private val FunctionName = "map_from_arrays" + private val ExpressionName = classOf[MapFromArrays].getSimpleName + + private case class Arm(name: String, comet: Boolean, dispatch: Boolean) + + private val Fallback = + Arm("Comet, dispatch off (Spark fallback)", comet = true, dispatch = false) + private val Dispatch = Arm("Comet, codegen dispatch", comet = true, dispatch = true) + private val Spark = Arm("Spark (Comet disabled)", comet = false, dispatch = false) + + private case class MapCase( + rows: Int, + arrayLength: Int, + duplicateKeys: Boolean, + mixed: Boolean) { + val name: String = { + val keys = if (duplicateKeys) "duplicate" else "unique" + val projection = if (mixed) "mixed projection" else "standalone" + s"$FunctionName: $keys keys, length $arrayLength, $projection -- $rows rows" + } + + val query: String = { + val neighbors = if (mixed) ", id + 1, length(label), substring(label, 1, 4)" else "" + s"SELECT map_from_arrays(k, v)$neighbors FROM parquetV1Table" + } + } + + private def configs(arm: Arm): Seq[(String, String)] = Seq( + SQLConf.MAP_KEY_DEDUP_POLICY.key -> "LAST_WIN", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> excludedRulesWith(ConstantFolding.ruleName), + CometConf.getExprEnabledConfigKey(ExpressionName) -> "true", + CometConf.getExprAllowIncompatConfigKey(ExpressionName) -> "false", + CometConf.COMET_ENABLED.key -> arm.comet.toString, + CometConf.COMET_EXEC_ENABLED.key -> arm.comet.toString, + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "true", + CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "true", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> arm.dispatch.toString) + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + require(mainArgs.length <= 1, "Usage: CometMapFromArraysBenchmark [input-element-budget]") + val elements = mainArgs.headOption.map(_.toInt).getOrElse(1024 * 1024) + require(elements >= 256, "The input-element budget must be at least 256") + require(spark.sparkContext.isLocal, "Runtime dispatcher verification requires local mode") + runBenchmark(s"$FunctionName: environment") { + emitEnvironment(elements) + } + for { + arrayLength <- Seq(4, 256) + duplicateKeys <- Seq(false, true) + } { + val rows = elements / arrayLength + withCorpus(rows, arrayLength, duplicateKeys) { + Seq(false, true).foreach { mixed => + val c = MapCase(rows, arrayLength, duplicateKeys, mixed) + runBenchmark(c.name) { + verifyArmsAgree(c) + val benchmark = new Benchmark(c.name, rows, output = output) + Seq(Fallback, Dispatch, Spark, Fallback.copy(name = Fallback.name + " (repeat)")) + .foreach { arm => + benchmark.addCase(arm.name) { _ => + withSQLConf(configs(arm): _*) { + spark.sql(c.query).noop() + } + } + } + benchmark.run() + } + } + } + } + } + + /** + * Compare maps independent of their Scala collection iteration order and rows as a multiset. + */ + private def verifyArmsAgree(c: MapCase): Unit = { + def collect(arm: Arm): Array[String] = { + // Spark 3.x's withSQLConf returns Unit; do not return a value from its block. + var rows: Array[Row] = Array.empty + withSQLConf(configs(arm): _*) { + CometScalaUDFCodegen.resetStats() + val df = spark.sql(c.query) + rows = df.collect() + checkPlan(c, arm, stripAQEPlan(df.queryExecution.executedPlan)) + val lookups = CometScalaUDFCodegen.stats().totalLookups + require( + (lookups > 0) == arm.dispatch, + s"${c.name}: ${arm.name} had $lookups dispatcher lookups; refusing to time it") + } + require(rows.length == c.rows, s"${c.name}: ${arm.name} produced ${rows.length} rows") + rows.map { row => + val entries = row.getMap[String, Long](0).toSeq.sortBy(_._1) + entries.mkString("[", ",", "]") + row.toSeq.drop(1).mkString("|", "|", "") + }.sorted + } + + val expected = collect(Spark) + Seq(Fallback, Dispatch).foreach { arm => + require( + expected.sameElements(collect(arm)), + s"${c.name}: ${arm.name} differs from Spark; refusing to time it") + } + emit("Verified all rows against Spark, expected plans and runtime dispatcher activity.") + } + + /** Abort instead of publishing a result under a misleading execution-path label. */ + private def checkPlan(c: MapCase, arm: Arm, plan: SparkPlan): Unit = { + val explain = new ExtendedExplainInfo() + val dispatched = explain.getCodegenDispatchExpressions(plan) + val native = explain.getNativeExpressions(plan) + val sparkMapProjection = plan.exists { + case p: ProjectExec => + p.projectList.exists(_.exists(_.isInstanceOf[MapFromArrays])) + case _ => false + } + val hasComet = plan.exists(_.nodeName.startsWith("Comet")) + val correctRoute = if (arm.dispatch) { + findFirstNonCometOperator(plan).isEmpty && + dispatched == Seq(FunctionName) && !native.contains(FunctionName) && !sparkMapProjection && + (!c.mixed || Seq("length", "substring").forall(native.contains)) + } else { + sparkMapProjection && dispatched.isEmpty && !native.contains(FunctionName) && + hasComet == arm.comet + } + require( + correctRoute, + s"${c.name}: unexpected ${arm.name} plan; refusing to time it.\n" + + s"Native expressions: $native; dispatched: $dispatched\n${plan.treeString}") + } + + private def withCorpus(rows: Int, arrayLength: Int, duplicateKeys: Boolean)( + f: => Unit): Unit = { + val distinctKeys = if (duplicateKeys) arrayLength / 2 else arrayLength + withTempPath { dir => + withTempTable("parquetV1Table") { + withSQLConf(configs(Spark): _*) { + val df = spark + .range(rows) + .selectExpr( + "id", + "concat('row-', cast(id as string)) as label", + s"transform(sequence(0, ${arrayLength - 1}), " + + s"x -> concat(cast(id as string), ':', cast(x % $distinctKeys as string))) as k", + s"transform(sequence(0, ${arrayLength - 1}), " + + s"x -> id * $arrayLength + x) as v") + prepareTable(dir, df) + } + f + } + } + } + + private def emitEnvironment(elements: Int): Unit = { + emit(s"Spark: ${spark.version}; Scala: ${scala.util.Properties.versionNumberString}") + emit(s"Java: ${System.getProperty("java.version")} (${System.getProperty("java.vm.name")})") + emit(s"OS: ${System.getProperty("os.name")} ${System.getProperty("os.arch")}") + emit(s"Master: ${spark.sparkContext.master}; JVM max heap: ${Runtime.getRuntime.maxMemory()}") + emit(s"Input-element budget: $elements; lengths: 4, 256; rows = budget / length") + emit( + "Duplicate keys: half as many distinct keys, each occurring twice with different values.") + emit("Preparation: Spark-written Snappy Parquet; each arm reads the same files.") + emit("Steady state: 2s warmup, at least 2 timed iterations and 2s timed execution per arm.") + emit( + "Compare the repeated fallback baseline and reported stdev before interpreting speedups.") + emit("Use the release build; these local timings do not establish production performance.") + emit(s"Batch size: ${CometConf.COMET_BATCH_SIZE.get(spark.sessionState.conf)}") + emit(s"Whole-stage codegen: ${spark.conf.get(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key)}") + configs(Dispatch).foreach { case (key, value) => emit(s"$key=$value") } + emit("Fallback: same settings, dispatcher disabled. Spark: Comet and dispatcher disabled.") + } + + private def emit(line: String): Unit = { + // scalastyle:off println + println(line) + // scalastyle:on println + output.foreach(_.write(s"$line\n".getBytes(StandardCharsets.UTF_8))) + } +}