From fd3ed02a5d797a7a5e1e24e39c97b3641f9e6118 Mon Sep 17 00:00:00 2001 From: Hung Date: Wed, 9 Sep 2026 21:38:01 +0800 Subject: [PATCH 1/8] feat: route MapSort fallback through the codegen dispatcher Keep native map_sort for scalar keys, and run Spark's MapSort.doGenCode in-pipeline for array, struct, and strict floating-point keys so the enclosing projection or shuffle does not fall back to Spark. --- .../expressions/_category_template/map.md | 13 +- .../shuffle/CometShuffleExchangeExec.scala | 7 +- .../org/apache/comet/serde/CometMapSort.scala | 3 +- .../apache/comet/CometCodegenAssertions.scala | 14 +- .../comet/CometMapExpressionSuite.scala | 130 +++++++++++++++++- .../exec/CometColumnarShuffleSuite.scala | 11 +- .../comet/exec/CometNativeShuffleSuite.scala | 37 +++-- 7 files changed, 185 insertions(+), 30 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/expressions/_category_template/map.md b/docs/source/user-guide/latest/compatibility/expressions/_category_template/map.md index 9368ae68d69..1676e4542da 100644 --- a/docs/source/user-guide/latest/compatibility/expressions/_category_template/map.md +++ b/docs/source/user-guide/latest/compatibility/expressions/_category_template/map.md @@ -23,12 +23,15 @@ under the License. Spark 4.0 inserts `MapSort` to normalize map values when they appear in shuffle hash partitioning keys, in `try_element_at`, and in other contexts where map ordering must be deterministic. Comet -runs `MapSort` natively, so map shuffle and group-by-on-map stay on Comet under Spark 4.0. +runs `MapSort` natively for supported scalar key types. Other orderable key types use Spark's own +generated code through the JVM codegen dispatcher, so the enclosing operator stays in the Comet +pipeline. -When `spark.comet.exec.strictFloatingPoint=true`, `MapSort` falls back to Spark for maps whose -keys contain `Float` or `Double` (consistent with `SortOrder` and `SortArray`). Arrow's sort uses -IEEE total ordering for floating-point, which differs from Spark's `Double.compare` semantics for -`NaN` and `-0.0`. +When `spark.comet.exec.strictFloatingPoint=true`, maps whose keys contain `Float` or `Double` also +use the dispatcher (consistent with `SortOrder` and `SortArray`). Arrow's sort uses IEEE total +ordering for floating-point, which differs from Spark's `Double.compare` semantics for `NaN` and +`-0.0`. If the dispatcher is disabled or cannot handle an expression, Comet safely falls back to +Spark. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 1f5d06a53af..4736e7188c6 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -452,9 +452,10 @@ object CometShuffleExchangeExec case MapType(keyType, valueType, _) if nestedHashPartitioningEnabled => // Map entry order is not semantically meaningful, so two equal maps must hash alike. // Spark 4.0+ normalizes a map shuffle key by wrapping it in `mapsort(...)`, which is - // gated separately by CometMapSort (scalar map keys only) and, when unsupported, fails - // the expression check below. Earlier Spark versions insert no such normalization, so - // Comet would hash physical entry order and could route equal maps differently. + // gated separately by CometMapSort. Scalar keys use native map_sort; other orderable key + // types can use Spark's generated MapSort code through the JVM dispatcher. Earlier Spark + // versions insert no such normalization, so Comet would hash physical entry order and + // could route equal maps differently. isSpark40Plus && supportedHashPartitioningDataType(keyType) && supportedHashPartitioningDataType(valueType) diff --git a/spark/src/main/spark-4.x/org/apache/comet/serde/CometMapSort.scala b/spark/src/main/spark-4.x/org/apache/comet/serde/CometMapSort.scala index 4fad289a080..f238ccd4068 100644 --- a/spark/src/main/spark-4.x/org/apache/comet/serde/CometMapSort.scala +++ b/spark/src/main/spark-4.x/org/apache/comet/serde/CometMapSort.scala @@ -25,7 +25,8 @@ import org.apache.spark.sql.types.MapType import org.apache.comet.CometConf import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, scalarFunctionExprToProtoWithReturnType, supportedScalarSortElementType} -object CometMapSort extends CometExpressionSerde[MapSort] { +// Key types unsupported by the native kernel can still run Spark's generated code in-pipeline. +object CometMapSort extends CometExpressionSerde[MapSort] with CodegenDispatchFallback { override def getIncompatibleReasons(): Seq[String] = Seq( diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenAssertions.scala b/spark/src/test/scala/org/apache/comet/CometCodegenAssertions.scala index bce8bfc5986..c77b39926fb 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenAssertions.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenAssertions.scala @@ -33,13 +33,23 @@ import org.apache.comet.vector.CometVector trait CometCodegenAssertions { /** Asserts the dispatcher actually ran during `f`, guarding against silent serde fallback. */ - protected def assertCodegenRan(f: => Unit): Unit = { + protected def assertCodegenRan[T](f: => T): T = { CometScalaUDFCodegen.resetStats() - f + val result = f val after = CometScalaUDFCodegen.stats() assert( after.compileCount + after.cacheHitCount >= 1, s"expected codegen dispatcher activity, got $after") + result + } + + /** Asserts the dispatcher did not run during `f`, guarding a native-path control case. */ + protected def assertCodegenDidNotRun[T](f: => T): T = { + CometScalaUDFCodegen.resetStats() + val result = f + val after = CometScalaUDFCodegen.stats() + assert(after.totalLookups == 0, s"expected no codegen dispatcher activity, got $after") + result } /** diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index ebbdce406a3..c38b8a56d69 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -22,7 +22,7 @@ package org.apache.comet import scala.util.Random import org.apache.hadoop.fs.Path -import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.{CometTestBase, DataFrame} import org.apache.spark.sql.catalyst.expressions.ArrayContains import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf @@ -31,7 +31,14 @@ 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 { + + private def assertMapSortInPlan(df: DataFrame): Unit = { + val plan = df.queryExecution.optimizedPlan + assert( + plan.exists(_.expressions.exists(_.exists(_.prettyName == "mapsort"))), + s"expected MapSort in optimized plan:\n$plan") + } test("read map[int, int] from parquet") { @@ -247,6 +254,125 @@ class CometMapExpressionSuite extends CometTestBase { } } + test("mapsort routes array keys through codegen dispatcher") { + assume(isSpark40Plus, "Spark 4.0 inserts MapSort for group-by on map keys") + withTable("t_map_sort_array_key") { + sql("CREATE TABLE t_map_sort_array_key (m MAP, INT>) USING parquet") + sql("""INSERT INTO t_map_sort_array_key VALUES + |(map(array(2, 1), 20, array(1, 2), 10)), + |(map(array(1, 2), 10, array(2, 1), 20)), + |(map(array(3), 30)), + |(NULL)""".stripMargin) + val df = sql("SELECT m, count(*) FROM t_map_sort_array_key GROUP BY m") + + assertMapSortInPlan(df) + val (_, cometPlan) = assertCodegenRan { + checkSparkAnswer(df) + } + val dispatched = new ExtendedExplainInfo().getCodegenDispatchExpressions(cometPlan) + assert( + dispatched.contains("mapsort"), + s"expected mapsort on codegen dispatch path, got $dispatched in:\n$cometPlan") + } + } + + test("mapsort routes struct keys through codegen dispatcher") { + assume(isSpark40Plus, "Spark 4.0 inserts MapSort for group-by on map keys") + withTable("t_map_sort_struct_key") { + sql("""CREATE TABLE t_map_sort_struct_key ( + | m MAP, INT>) USING parquet""".stripMargin) + sql("""INSERT INTO t_map_sort_struct_key VALUES + |(map(named_struct('a', 2, 'b', 'b'), 20, + | named_struct('a', 1, 'b', 'a'), 10)), + |(map(named_struct('a', 1, 'b', 'a'), 10, + | named_struct('a', 2, 'b', 'b'), 20)), + |(map(named_struct('a', 3, 'b', 'c'), 30)), + |(NULL)""".stripMargin) + val df = sql("SELECT m, count(*) FROM t_map_sort_struct_key GROUP BY m") + + assertMapSortInPlan(df) + val (_, cometPlan) = assertCodegenRan { + checkSparkAnswer(df) + } + val dispatched = new ExtendedExplainInfo().getCodegenDispatchExpressions(cometPlan) + assert( + dispatched.contains("mapsort"), + s"expected mapsort on codegen dispatch path, got $dispatched in:\n$cometPlan") + } + } + + test("mapsort keeps scalar keys on the native path") { + assume(isSpark40Plus, "Spark 4.0 inserts MapSort for group-by on map keys") + withTable("t_map_sort_scalar_key") { + sql("CREATE TABLE t_map_sort_scalar_key (m MAP) USING parquet") + sql("""INSERT INTO t_map_sort_scalar_key VALUES + |(map(2, 20, 1, 10)), + |(map(1, 10, 2, 20)), + |(map(3, 30)), + |(NULL)""".stripMargin) + val df = sql("SELECT m, count(*) FROM t_map_sort_scalar_key GROUP BY m") + + assertMapSortInPlan(df) + val (_, cometPlan) = assertCodegenDidNotRun(checkSparkAnswer(df)) + val explain = new ExtendedExplainInfo() + val nativeExpressions = explain.getNativeExpressions(cometPlan) + assert( + nativeExpressions.contains("mapsort"), + s"expected native mapsort expression, got $nativeExpressions in:\n$cometPlan") + assert( + !explain.getCodegenDispatchExpressions(cometPlan).contains("mapsort"), + s"scalar-key mapsort should not use codegen dispatch:\n$cometPlan") + } + } + + test("mapsort with complex keys falls back when codegen dispatcher is disabled") { + assume(isSpark40Plus, "Spark 4.0 inserts MapSort for group-by on map keys") + withTable("t_map_sort_dispatch_disabled") { + sql("CREATE TABLE t_map_sort_dispatch_disabled (m MAP, INT>) USING parquet") + sql("""INSERT INTO t_map_sort_dispatch_disabled VALUES + |(map(array(2, 1), 20, array(1, 2), 10)), + |(map(array(1, 2), 10, array(2, 1), 20))""".stripMargin) + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") { + val df = sql("SELECT m, count(*) FROM t_map_sort_dispatch_disabled GROUP BY m") + + assertMapSortInPlan(df) + assertCodegenDidNotRun { + checkSparkAnswerAndFallbackReason( + df, + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key + "=false") + } + } + } + } + + // Spark rejects collated strings as map keys (`UNSUPPORTED_FEATURE.COLLATIONS_IN_MAP_KEYS`), so + // `MapSort` never sees that shape. `supportedScalarSortElementType` still excludes them, and the + // same `Unsupported` → dispatcher path is covered by the array/struct cases above. + + test("mapsort routes strict floating-point keys through codegen dispatcher") { + assume(isSpark40Plus, "Spark 4.0 inserts MapSort for group-by on map keys") + withSQLConf(CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") { + withTable("t_map_sort_fp_key") { + sql("CREATE TABLE t_map_sort_fp_key (m MAP) USING parquet") + sql("""INSERT INTO t_map_sort_fp_key VALUES + |(map(CAST('NaN' AS DOUBLE), 1, CAST('-0.0' AS DOUBLE), 2, 1.0, 3)), + |(map(1.0, 3, CAST('-0.0' AS DOUBLE), 2, CAST('NaN' AS DOUBLE), 1)), + |(map(0.0, 4)), + |(NULL)""".stripMargin) + val df = sql("SELECT m, count(*) FROM t_map_sort_fp_key GROUP BY m") + + assertMapSortInPlan(df) + val (_, cometPlan) = assertCodegenRan { + checkSparkAnswer(df) + } + val dispatched = new ExtendedExplainInfo().getCodegenDispatchExpressions(cometPlan) + assert( + dispatched.contains("mapsort"), + s"expected mapsort on codegen dispatch path, got $dispatched in:\n$cometPlan") + } + } + } + test("map_from_entries - binary type routes through codegen dispatcher") { val table = "t2" withTable(table) { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala index bf354356386..9afff38d417 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala @@ -38,7 +38,6 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper { protected val adaptiveExecutionEnabled: Boolean @@ -142,10 +141,8 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar } test("columnar shuffle on array/struct map key/value") { - // Spark 4.0 normalizes maps used as shuffle keys with mapsort(...). Comet's map_sort - // relies on Arrow's sort_to_indices, which only supports scalar key types, so a map - // with array or struct keys cannot be sorted natively and the shuffle falls back. - val complexKeyShuffles = if (isSpark40Plus) 0 else 1 + // Spark 4.0+ normalizes maps used as shuffle keys with mapsort(...). Native map_sort only + // supports scalar keys, so array and struct keys use Spark codegen in the Comet pipeline. Seq("false", "true").foreach { execEnabled => Seq(10, 201).foreach { numPartitions => Seq("1.0", "10.0").foreach { ratio => @@ -158,7 +155,7 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar .repartition(numPartitions, $"_1", $"_2") .sortWithinPartitions($"_2") - checkShuffleAnswer(df, complexKeyShuffles) + checkShuffleAnswer(df, 1) } withParquetTable((0 until 50).map(i => (Map(i -> Seq(i, i + 1)), i + 1)), "tbl") { @@ -176,7 +173,7 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar .repartition(numPartitions, $"_1", $"_2") .sortWithinPartitions($"_2") - checkShuffleAnswer(df, complexKeyShuffles) + checkShuffleAnswer(df, 1) } withParquetTable((0 until 50).map(i => (Map(i -> ((i, i.toString))), i + 1)), "tbl") { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index ff6ef3b10f0..b0be7856038 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -49,12 +49,15 @@ import org.apache.spark.sql.functions.{col, count, sum} import org.apache.spark.sql.types.{ArrayType, DataType, LongType, MapType, StructField, StructType} import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} -import org.apache.comet.{CometConf, CometExecIterator, CometExplainInfo, CometShuffleBlockIterator, CometShuffleSizeLimitException, Native} +import org.apache.comet.{CometCodegenAssertions, CometConf, CometExecIterator, CometExplainInfo, CometShuffleBlockIterator, CometShuffleSizeLimitException, Native} import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus import org.apache.comet.serde.{OperatorOuterClass, PartitioningOuterClass} import org.apache.comet.shuffle.ShufflePartitionPusher -class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper { +class CometNativeShuffleSuite + extends CometTestBase + with AdaptiveSparkPlanHelper + with CometCodegenAssertions { override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit pos: Position): Unit = { super.test(testName, testTags: _*) { @@ -878,7 +881,8 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper // Map entry order carries no meaning, so equal maps must hash alike. Spark 4.0+ normalizes a // map shuffle key with `mapsort(...)`; earlier versions do not, so Comet must not hash a raw // map there. The gate therefore only admits map keys on Spark 4.0+, and only when the - // `mapsort` itself is convertible (CometMapSort supports scalar map keys only). + // `mapsort` itself can stay in the Comet pipeline, either natively or through the JVM + // codegen dispatcher. withParquetTable((0 until 50).map(i => (i, Map(i % 7 -> (i % 5)))), "tbl") { val df = sql("SELECT * FROM tbl").repartition(10, $"_2").sortWithinPartitions($"_1") @@ -887,15 +891,28 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper } } - test("native shuffle on map hash partitioning key with non-scalar map key falls back") { - // A map whose own key is nested cannot be `mapsort`ed by Comet (Arrow's sort_to_indices - // handles scalar keys only), so the normalization Spark 4.0+ requires is unavailable and the - // shuffle must fall back rather than hash an unnormalized map. + test("native shuffle on map hash partitioning key with non-scalar map key uses dispatcher") { + // Arrow's sort_to_indices only handles scalar keys, so Spark's MapSort.doGenCode performs the + // normalization through the JVM codegen dispatcher while the shuffle remains native. assume(isSpark40Plus, "map shuffle keys are only normalized with mapsort on Spark 4.0+") - withParquetTable((0 until 50).map(i => (i, Map(Seq(i % 7) -> (i % 5)))), "tbl") { - val df = sql("SELECT * FROM tbl").repartition(10, $"_2").sortWithinPartitions($"_1") + withNestedHashPartitioning { + withParquetTable((0 until 50).map(i => (i, Map(Seq(i % 7) -> (i % 5)))), "tbl") { + val df = sql("SELECT * FROM tbl").repartition(10, $"_2").sortWithinPartitions($"_1") - checkShuffleAnswer(df, 0) + assertCodegenRan { + checkShuffleAnswer(df, 1) + } + } + + withParquetTable( + (0 until 50).map(i => (i, Map((i % 7, (i % 5).toString) -> (i % 3)))), + "tbl") { + val df = sql("SELECT * FROM tbl").repartition(10, $"_2").sortWithinPartitions($"_1") + + assertCodegenRan { + checkShuffleAnswer(df, 1) + } + } } } From 0fa0f9c09927d77e2b6a90565d0ee22663ba112b Mon Sep 17 00:00:00 2001 From: Hung Date: Thu, 10 Sep 2026 14:42:34 +0800 Subject: [PATCH 2/8] test: add a matched MapSort dispatcher vs Spark-fallback microbenchmark Compare in-pipeline MapSort.doGenCode against the previous whole-operator fallback for array, struct, and strict floating-point keys, so staying in Comet is measured rather than assumed. --- .../sql/benchmark/CometMapSortBenchmark.scala | 661 ++++++++++++++++++ 1 file changed, 661 insertions(+) create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala new file mode 100644 index 00000000000..4d1c5151cd7 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala @@ -0,0 +1,661 @@ +/* + * 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.{DataFrame, Row} +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec +import org.apache.spark.sql.execution.ProjectExec +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.functions.{col, spark_partition_id} +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.{CometConf, ExtendedExplainInfo} +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus +import org.apache.comet.udf.codegen.CometScalaUDFCodegen + +/** + * Matched benchmark for the two routes Spark 4.x can take for `MapSort` shapes that Comet cannot + * sort natively: + * + * - with the JVM codegen dispatcher disabled, the enclosing projection or shuffle falls back to + * Spark; and + * - with it enabled, Spark's `MapSort.doGenCode` executes inside the Comet pipeline. + * + * Every pair reads the same Parquet data and differs only in + * `spark.comet.exec.scalaUDF.codegen.enabled`. Array and struct cases vary map size independently + * from nested-key width; strict floating-point cases include NaN and signed zero. Input maps are + * written in reverse key order and one row in 64 has a NULL map. + * + * Spark 4.0 and 4.1 only insert `MapSort` for grouping and repartition expressions; + * `try_element_at` itself does not insert one. To measure a projection without also timing an + * aggregate, `mapSortProjection` asks Spark's grouping optimizer to construct its real + * `Project(MapSort(m))`, then executes that logical Project on its own. This avoids importing the + * Spark-4.x-only `MapSort` class and keeps this common benchmark source compilable on Spark 3.x. + * + * To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make \ + * benchmark-org.apache.spark.sql.benchmark.CometMapSortBenchmark + * }}} + * + * Formal steady-state results require at least five fresh JVM invocations, alternating + * `-Dcomet.mapSortBenchmark.caseOrder=fallback-first` and `dispatcher-first`, and reporting the + * median/min/max across processes rather than only Benchmark's within-process best time. + * + * First-action latency is a separate invocation mode. Each process measures one action only, so + * callers should run every schema/route/workload tuple at least three times in fresh JVMs: + * {{{ + * -Dcomet.mapSortBenchmark.mode=first-action \ + * -Dcomet.mapSortBenchmark.shape=array-small \ + * -Dcomet.mapSortBenchmark.route=dispatcher \ + * -Dcomet.mapSortBenchmark.workload=projection \ + * -Dcomet.mapSortBenchmark.repetition=1 + * }}} + */ +object CometMapSortBenchmark extends CometBenchmarkBase { + + private val DefaultProjectionRows = 1000000 + private val DefaultShuffleRows = 250000 + private val DefaultInputPartitions = 4 + private val DefaultShufflePartitions = 16 + private val DefaultVerificationRows = 2048 + private val DefaultFirstActionRows = 1024 + private val NullMapEvery = 64 + + private val ModeProperty = "comet.mapSortBenchmark.mode" + private val CaseOrderProperty = "comet.mapSortBenchmark.caseOrder" + private val ShapeProperty = "comet.mapSortBenchmark.shape" + private val RouteProperty = "comet.mapSortBenchmark.route" + private val WorkloadProperty = "comet.mapSortBenchmark.workload" + private val RepetitionProperty = "comet.mapSortBenchmark.repetition" + + private val Mode = sys.props.getOrElse(ModeProperty, "steady") + private val CaseOrder = sys.props.getOrElse(CaseOrderProperty, "fallback-first") + + // These overrides make plan/routing validation practical on a development machine. The normal + // microbenchmark runner supplies none of them and therefore always uses the values above. The + // effective values are emitted into the results file. + private val ProjectionRows = + intProperty("comet.mapSortBenchmark.projectionRows", DefaultProjectionRows) + private val ShuffleRows = intProperty("comet.mapSortBenchmark.shuffleRows", DefaultShuffleRows) + private val InputPartitions = + intProperty("comet.mapSortBenchmark.inputPartitions", DefaultInputPartitions) + private val ShufflePartitions = + intProperty("comet.mapSortBenchmark.shufflePartitions", DefaultShufflePartitions) + private val VerificationRows = + intProperty("comet.mapSortBenchmark.verificationRows", DefaultVerificationRows) + private val FirstActionRows = + intProperty("comet.mapSortBenchmark.firstActionRows", DefaultFirstActionRows) + + private sealed trait KeyFamily { + def label: String + def mapType(width: Int): String + def key(entry: String, width: Int): String + } + + private case object ArrayKey extends KeyFamily { + override val label: String = "array" + + override def mapType(width: Int): String = "MAP, INT>" + + override def key(entry: String, width: Int): String = + s"""transform( + | sequence(0, ${width - 1}), + | j -> CAST(pmod(id, 1000003) * 4096 + CAST($entry AS BIGINT) * $width + j AS INT)) + |""".stripMargin.replace('\n', ' ') + } + + private case object StructKey extends KeyFamily { + override val label: String = "struct" + + override def mapType(width: Int): String = { + val fields = (0 until width).map(i => s"f$i: INT").mkString(", ") + s"MAP, INT>" + } + + override def key(entry: String, width: Int): String = { + val fields = (0 until width).flatMap { i => + Seq( + s"'f$i'", + s"CAST(pmod(id, 1000003) * 4096 + CAST($entry AS BIGINT) * $width + $i AS INT)") + } + s"named_struct(${fields.mkString(", ")})" + } + } + + private case object StrictDoubleKey extends KeyFamily { + override val label: String = "strict-double" + + override def mapType(width: Int): String = "MAP" + + override def key(entry: String, width: Int): String = + s"""CASE + | WHEN $entry = 0 THEN CAST('NaN' AS DOUBLE) + | WHEN $entry = 1 THEN CAST('-0.0' AS DOUBLE) + | ELSE CAST(id * 128 + CAST($entry AS BIGINT) + 1 AS DOUBLE) + |END""".stripMargin.replace('\n', ' ') + } + + private case class Shape(name: String, family: KeyFamily, mapSize: Int, keyWidth: Int) { + require(mapSize >= 2, "mapSize must leave room for NaN and -0.0") + require(keyWidth > 0, "keyWidth must be positive") + + def description: String = + if (family == StrictDoubleKey) { + s"${family.label}, map-size=$mapSize" + } else { + s"${family.label}, map-size=$mapSize, key-width=$keyWidth" + } + } + + private val ArraySmall = Shape("array-small", ArrayKey, mapSize = 4, keyWidth = 2) + private val StructNarrow = Shape("struct-small", StructKey, mapSize = 4, keyWidth = 2) + private val StructWide = Shape("struct-wide", StructKey, mapSize = 4, keyWidth = 8) + + private val Shapes = Seq( + ArraySmall, + Shape("array-large-map", ArrayKey, mapSize = 32, keyWidth = 2), + Shape("array-wide", ArrayKey, mapSize = 4, keyWidth = 8), + StructNarrow, + Shape("struct-large-map", StructKey, mapSize = 32, keyWidth = 2), + StructWide, + Shape("double-small", StrictDoubleKey, mapSize = 4, keyWidth = 1), + Shape("double-large-map", StrictDoubleKey, mapSize = 32, keyWidth = 1)) + + private val FallbackCaseName = "Comet / Spark fallback" + private val DispatcherCaseName = "Comet / MapSort dispatcher" + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + runBenchmark("MapSort dispatcher: environment") { + emitEnvironment() + } + + if (!isSpark40Plus) { + emit(s"SKIPPED: Spark ${spark.version} does not define or insert MapSort (requires 4.0+).") + return + } + + Mode match { + case "first-action" => + runFirstAction() + return + case "steady" => + case other => + throw new IllegalArgumentException( + s"invalid $ModeProperty=$other (expected steady or first-action)") + } + + Shapes.foreach(runProjectionBenchmark) + Shapes.foreach(runShuffleBenchmark) + } + + private def runProjectionBenchmark(shape: Shape): Unit = { + withCorpus(shape, ProjectionRows) { + verifyMatchedPair(shape, "projection", () => mapSortProjection(VerificationRows)) + // Plan the dispatcher arm first and use distinct logical trees. Comet records planning + // diagnostics in TreeNode tags; reusing (or first fallback-tagging) the same Catalyst tree + // can otherwise make the second arm appear to have inherited the first arm's route. + val dispatcher = prepareQuery( + shape, + dispatch = true, + shuffle = false, + mapSortProjection().queryExecution.logical) + val fallback = prepareQuery( + shape, + dispatch = false, + shuffle = false, + mapSortProjection().queryExecution.logical) + assertRoute( + shape, + "projection", + dispatch = true, + stripAQEPlan(dispatcher.queryExecution.executedPlan)) + assertRoute( + shape, + "projection", + dispatch = false, + stripAQEPlan(fallback.queryExecution.executedPlan)) + runBenchmark(s"MapSort projection -- ${shape.description}") { + val benchmark = new Benchmark( + s"MapSort projection -- ${shape.description}", + ProjectionRows, + output = output) + addMatchedCases(benchmark, shape, shuffle = false, fallback, dispatcher) + benchmark.run() + } + } + } + + private def runShuffleBenchmark(shape: Shape): Unit = { + withCorpus(shape, ShuffleRows) { + verifyMatchedPair(shape, "shuffle", () => shuffleQuery(VerificationRows)) + val dispatcher = prepareQuery( + shape, + dispatch = true, + shuffle = true, + shuffleQuery().queryExecution.logical) + val fallback = prepareQuery( + shape, + dispatch = false, + shuffle = true, + shuffleQuery().queryExecution.logical) + assertRoute( + shape, + "shuffle", + dispatch = true, + stripAQEPlan(dispatcher.queryExecution.executedPlan)) + assertRoute( + shape, + "shuffle", + dispatch = false, + stripAQEPlan(fallback.queryExecution.executedPlan)) + runBenchmark(s"MapSort native shuffle -- ${shape.description}") { + val benchmark = new Benchmark( + s"MapSort native shuffle -- ${shape.description}", + ShuffleRows, + output = output) + addMatchedCases(benchmark, shape, shuffle = true, fallback, dispatcher) + benchmark.run() + } + } + } + + /** + * Builds the exact MapSort expression inserted by Spark's grouping optimizer, but returns only + * that projection. The aggregate is a construction device and is never part of the returned + * DataFrame or the timed execution. + */ + private def mapSortProjection(maxRows: Int = Int.MaxValue): DataFrame = { + val input = limitedInput(maxRows) + val optimizedGrouping = input.groupBy(col("m")).count().queryExecution.optimizedPlan + val mapSortProject = optimizedGrouping + .collectFirst { + case plan + if plan.output.exists(_.name == "_groupingmapsort") && + plan.expressions.exists(containsMapSort) => + plan + } + .getOrElse { + throw new IllegalStateException( + "Spark did not insert the expected MapSort grouping projection:\n" + + optimizedGrouping.treeString) + } + + val projected = dataFrameOfRows(mapSortProject) + .select(col("_groupingmapsort").as("sorted_m")) + assertMapSortInOptimizedPlan(projected) + projected + } + + /** + * Spark 4.0 moved the Dataset implementation and its `ofRows` factory to `sql.classic`, while + * Spark 3.x keeps it in `sql`. Reflection across that packaging-only difference lets the common + * source compile on every supported Spark line. + */ + private def dataFrameOfRows(plan: LogicalPlan): DataFrame = { + val companionClass = + Seq("org.apache.spark.sql.classic.Dataset$", "org.apache.spark.sql.Dataset$").iterator + .map(name => scala.util.Try(Class.forName(name)).toOption) + .collectFirst { case Some(clazz) => clazz } + .getOrElse(throw new IllegalStateException("could not locate Spark Dataset companion")) + val module = companionClass.getField("MODULE$").get(null) + val ofRows = companionClass.getMethods + .find(method => method.getName == "ofRows" && method.getParameterCount == 2) + .getOrElse(throw new IllegalStateException("could not locate Spark Dataset.ofRows")) + ofRows.invoke(module, spark, plan).asInstanceOf[DataFrame] + } + + private def shuffleQuery(maxRows: Int = Int.MaxValue): DataFrame = { + val shuffled = limitedInput(maxRows).repartition(ShufflePartitions, col("m")) + assertMapSortInOptimizedPlan(shuffled) + shuffled + } + + private def limitedInput(maxRows: Int): DataFrame = { + val input = spark.table("parquetV1Table") + if (maxRows == Int.MaxValue) input else input.where(col("id") < maxRows) + } + + private def containsMapSort( + expression: org.apache.spark.sql.catalyst.expressions.Expression): Boolean = + expression.exists(_.prettyName == "mapsort") + + private def assertMapSortInOptimizedPlan(df: DataFrame): Unit = { + val plan = df.queryExecution.optimizedPlan + assert( + plan.exists(_.expressions.exists(containsMapSort)), + s"expected MapSort in optimized plan:\n${plan.treeString}") + } + + /** + * Executes both routes on a small prefix of the same Parquet corpus and checks results/plans. + */ + private def verifyMatchedPair(shape: Shape, workload: String, query: () => DataFrame): Unit = { + // See runProjectionBenchmark: route the dispatcher tree before adding any fallback tags. + val dispatcher = captureRun(shape, dispatch = true, workload, query) + val fallback = captureRun(shape, dispatch = false, workload, query) + + assert( + fallback.rows.sameElements(dispatcher.rows), + s"${shape.description} $workload routes produced different rows") + val explain = new ExtendedExplainInfo() + assert( + !explain.getCodegenDispatchExpressions(fallback.plan).contains("mapsort"), + s"MapSort was unexpectedly annotated as dispatched in $FallbackCaseName:\n" + + fallback.plan.treeString) + assert( + explain.getCodegenDispatchExpressions(dispatcher.plan).contains("mapsort"), + s"MapSort was not annotated as dispatched for ${shape.description} $workload:\n" + + dispatcher.plan.treeString) + + workload match { + case "projection" => + assert( + dispatcher.firstNonComet.isEmpty, + s"dispatcher projection was not fully Comet: ${dispatcher.firstNonComet}\n" + + dispatcher.plan.treeString) + assert( + fallback.plan.exists(_.isInstanceOf[ProjectExec]), + s"fallback route did not contain a Spark ProjectExec:\n${fallback.plan.treeString}") + + case "shuffle" => + assert( + dispatcher.plan.exists(_.isInstanceOf[CometShuffleExchangeExec]), + s"dispatcher route did not retain Comet native shuffle:\n${dispatcher.plan.treeString}") + assert( + !fallback.plan.exists(_.isInstanceOf[CometShuffleExchangeExec]) && + fallback.plan.exists(_.isInstanceOf[ShuffleExchangeExec]), + s"fallback route did not use Spark shuffle exclusively:\n${fallback.plan.treeString}") + + case other => throw new IllegalArgumentException(s"unknown workload: $other") + } + + val equality = + if (workload == "shuffle") "equal rows and spark_partition_id assignments" + else "equal results" + emit( + s"Verified ${shape.description} $workload on ${dispatcher.rows.length} rows: " + + s"$equality; $FallbackCaseName used Spark; $DispatcherCaseName dispatched mapsort.") + emit(s" fallback executed plan: ${oneLine(fallback.plan.treeString)}") + emit(s" dispatcher executed plan: ${oneLine(dispatcher.plan.treeString)}") + emit( + " dispatcher codegen expressions: " + + new ExtendedExplainInfo().getCodegenDispatchExpressions(dispatcher.plan)) + } + + private case class CapturedRun( + rows: Array[String], + plan: org.apache.spark.sql.execution.SparkPlan, + firstNonComet: Option[String]) + + private def assertRoute( + shape: Shape, + workload: String, + dispatch: Boolean, + plan: org.apache.spark.sql.execution.SparkPlan): Unit = { + val mapSortDispatched = + new ExtendedExplainInfo().getCodegenDispatchExpressions(plan).contains("mapsort") + assert( + mapSortDispatched == dispatch, + s"unexpected MapSort dispatch annotation for ${shape.description} $workload, " + + s"dispatch=$dispatch:\n${plan.treeString}") + (workload, dispatch) match { + case ("projection", true) => + assert( + findFirstNonCometOperator(plan).isEmpty, + s"dispatcher projection was not fully Comet:\n${plan.treeString}") + case ("projection", false) => + assert( + plan.exists(_.isInstanceOf[ProjectExec]), + s"fallback route did not contain a Spark ProjectExec:\n${plan.treeString}") + case ("shuffle", true) => + assert( + plan.exists(_.isInstanceOf[CometShuffleExchangeExec]), + s"dispatcher route did not retain Comet native shuffle:\n${plan.treeString}") + case ("shuffle", false) => + assert( + !plan.exists(_.isInstanceOf[CometShuffleExchangeExec]) && + plan.exists(_.isInstanceOf[ShuffleExchangeExec]), + s"fallback route did not use Spark shuffle exclusively:\n${plan.treeString}") + case _ => + throw new IllegalArgumentException(s"unknown workload: $workload") + } + } + + private def captureRun( + shape: Shape, + dispatch: Boolean, + workload: String, + query: () => DataFrame): CapturedRun = { + var result: CapturedRun = null + withSQLConf(configs(shape, dispatch, workload == "shuffle"): _*) { + val df = query() + val checked = + if (workload == "shuffle") { + // Row equality alone cannot detect a different hash-partition assignment. Preserve the + // entire row and append the partition id so the comparison validates both. + df.select(col("*"), spark_partition_id().as("_partition_id")) + } else { + df + } + val rows = checked.collect().map(renderRow).sorted + val plan = stripAQEPlan(df.queryExecution.executedPlan) + result = CapturedRun(rows, plan, findFirstNonCometOperator(plan).map(_.nodeName)) + } + result + } + + private def renderRow(row: Row): String = row.toSeq.map(String.valueOf).mkString("|") + + private def oneLine(value: String): String = + value.split("\\n").iterator.map(_.trim).mkString(" | ") + + /** Materializes physical planning outside Benchmark's timed closure. */ + private def prepareQuery( + shape: Shape, + dispatch: Boolean, + shuffle: Boolean, + logicalPlan: LogicalPlan): DataFrame = { + var prepared: DataFrame = null + withSQLConf(configs(shape, dispatch, shuffle): _*) { + prepared = dataFrameOfRows(logicalPlan) + prepared.queryExecution.executedPlan + } + prepared + } + + private def runPreparedQuery( + shape: Shape, + dispatch: Boolean, + shuffle: Boolean, + df: DataFrame): Unit = + withSQLConf(configs(shape, dispatch, shuffle): _*) { + df.noop() + } + + private def addMatchedCases( + benchmark: Benchmark, + shape: Shape, + shuffle: Boolean, + fallback: DataFrame, + dispatcher: DataFrame): Unit = { + def addFallback(): Unit = benchmark.addCase(FallbackCaseName) { _ => + runPreparedQuery(shape, dispatch = false, shuffle = shuffle, df = fallback) + } + def addDispatcher(): Unit = benchmark.addCase(DispatcherCaseName) { _ => + runPreparedQuery(shape, dispatch = true, shuffle = shuffle, df = dispatcher) + } + + CaseOrder match { + case "fallback-first" => + addFallback() + addDispatcher() + case "dispatcher-first" => + addDispatcher() + addFallback() + case other => + throw new IllegalArgumentException( + s"invalid $CaseOrderProperty=$other (expected fallback-first or dispatcher-first)") + } + } + + private def configs(shape: Shape, dispatch: Boolean, shuffle: Boolean): Seq[(String, String)] = + Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> dispatch.toString, + CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> + (shape.family == StrictDoubleKey).toString, + CometConf.getExprAllowIncompatConfigKey("MapSort") -> "false", + CometConf.COMET_SHUFFLE_ENABLED.key -> shuffle.toString, + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_NESTED_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> ShufflePartitions.toString) + + /** + * Measures exactly one first action for one route/workload/schema. A valid first-action study + * invokes this mode in a fresh JVM for every sample; the benchmark deliberately does not call a + * second action or claim that the dispatcher counters prove how much Java compilation occurred. + */ + private def runFirstAction(): Unit = { + val shapeName = requiredProperty(ShapeProperty) + val shape = Shapes + .find(_.name == shapeName) + .getOrElse(throw new IllegalArgumentException( + s"invalid $ShapeProperty=$shapeName; expected one of ${Shapes.map(_.name).mkString(",")}")) + val workload = requiredProperty(WorkloadProperty) + require( + workload == "projection" || workload == "shuffle", + s"$WorkloadProperty must be projection or shuffle") + val route = requiredProperty(RouteProperty) + require(route == "fallback" || route == "dispatcher", s"$RouteProperty is invalid: $route") + val dispatch = route == "dispatcher" + val shuffle = workload == "shuffle" + + withCorpus(shape, FirstActionRows) { + val logicalPlan = + if (shuffle) shuffleQuery().queryExecution.logical + else mapSortProjection().queryExecution.logical + val df = prepareQuery(shape, dispatch, shuffle, logicalPlan) + CometScalaUDFCodegen.resetStats() + val elapsed = + timeMillis(runPreparedQuery(shape, dispatch = dispatch, shuffle = shuffle, df = df)) + val plan = stripAQEPlan(df.queryExecution.executedPlan) + val stats = CometScalaUDFCodegen.stats() + assertRoute(shape, workload, dispatch, plan) + val repetition = sys.props.getOrElse(RepetitionProperty, "") + emit( + f"MAPSORT_FIRST_ACTION shape=${shape.name} route=$route workload=$workload " + + f"repetition=$repetition rows=$FirstActionRows elapsed_ms=$elapsed%.1f " + + s"dispatcher_compile_count=${stats.compileCount} " + + s"dispatcher_cache_hit_count=${stats.cacheHitCount}") + emit(s" executed plan: ${oneLine(plan.treeString)}") + emit( + "This is first-action latency from one process. Dispatcher counters are routing/cache " + + "observations, not a measurement of Java compiler time.") + } + } + + private def timeMillis(f: => Unit): Double = { + val start = System.nanoTime() + f + (System.nanoTime() - start) / 1e6 + } + + private def withCorpus(shape: Shape, rows: Int)(f: => Unit): Unit = { + withTempPath { dir => + withTempTable(tbl, "parquetV1Table") { + spark + .range(0L, rows.toLong, 1L, InputPartitions) + .createOrReplaceTempView(tbl) + prepareTable(dir, spark.sql(corpusQuery(shape))) + f + } + } + } + + private def corpusQuery(shape: Shape): String = { + val entrySequence = s"sequence(${shape.mapSize - 1}, 0, -1)" + val keys = + s"transform($entrySequence, i -> ${shape.family.key("i", shape.keyWidth)})" + val values = + s"transform($entrySequence, i -> CAST(pmod(id * 17 + i, 2147483647) AS INT))" + val map = s"map_from_arrays($keys, $values)" + val lookupEntry = s"CAST(pmod(id, ${shape.mapSize}) AS INT)" + val lookupKey = shape.family.key(lookupEntry, shape.keyWidth) + + s""" + |SELECT + | id, + | CASE WHEN pmod(id, $NullMapEvery) = 0 + | THEN CAST(NULL AS ${shape.family.mapType(shape.keyWidth)}) + | ELSE $map + | END AS m, + | $lookupKey AS lookup_key + |FROM $tbl + |""".stripMargin + } + + private def emitEnvironment(): Unit = { + emit(s"Spark version: ${spark.version}") + emit( + s"Java version: ${System.getProperty("java.version")} " + + s"(${System.getProperty("java.vm.name")})") + emit(s"Scala version: ${scala.util.Properties.versionNumberString}") + emit(s"spark.master: ${spark.conf.get("spark.master", "")}") + emit( + s"${CometConf.COMET_BATCH_SIZE.key}: " + + CometConf.COMET_BATCH_SIZE.get(spark.sessionState.conf)) + emit(s"Projection rows: $ProjectionRows; shuffle rows: $ShuffleRows") + emit(s"Input partitions: $InputPartitions; shuffle partitions: $ShufflePartitions") + emit(s"Correctness/routing prefix rows: $VerificationRows; NULL map density: 1/$NullMapEvery") + emit(s"Mode: $Mode; steady-state case order: $CaseOrder") + emit("Shuffle mode: native; AQE: disabled; nested hash partitioning: enabled") + emit(s"Only matched-arm difference: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}") + emit(s"Pinned: ${CometConf.getExprAllowIncompatConfigKey("MapSort")}=false") + emit("Projection note: Spark 4.0/4.1 do not insert MapSort for try_element_at. This suite") + emit(" executes the optimizer-inserted grouping MapSort Project alone; no aggregate is run.") + emit(s"Shapes: ${Shapes.map(_.description).mkString("; ")}") + } + + /** + * [[Benchmark]] tees its output; custom environment/first-use tables need the same behaviour. + */ + private def emit(line: String): Unit = { + // scalastyle:off println + println(line) + // scalastyle:on println + output.foreach(_.write(s"$line\n".getBytes(StandardCharsets.UTF_8))) + } + + private def intProperty(name: String, default: Int): Int = { + val value = sys.props.get(name).map(_.toInt).getOrElse(default) + require(value > 0, s"$name must be positive") + value + } + + private def requiredProperty(name: String): String = + sys.props.getOrElse(name, throw new IllegalArgumentException(s"missing required -D$name")) +} From f1a7c2e641ad26b74f86bc6bbe09d81bc58317ef Mon Sep 17 00:00:00 2001 From: Hung Date: Fri, 11 Sep 2026 23:35:42 +0800 Subject: [PATCH 3/8] test: time the prepared MapSort physical plan instead of a noop write DatasetToBenchmark.noop() builds a separate write QueryExecution on Spark 4.x, so route validation and the timed action were not the same plan. --- .../sql/benchmark/CometMapSortBenchmark.scala | 204 ++++++++++++++---- 1 file changed, 163 insertions(+), 41 deletions(-) diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala index 4d1c5151cd7..591f8596230 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala @@ -25,7 +25,7 @@ import org.apache.spark.benchmark.Benchmark import org.apache.spark.sql.{DataFrame, Row} import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec -import org.apache.spark.sql.execution.ProjectExec +import org.apache.spark.sql.execution.{ProjectExec, SparkPlan, SQLExecution} import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.functions.{col, spark_partition_id} import org.apache.spark.sql.internal.SQLConf @@ -53,6 +53,12 @@ import org.apache.comet.udf.codegen.CometScalaUDFCodegen * `Project(MapSort(m))`, then executes that logical Project on its own. This avoids importing the * Spark-4.x-only `MapSort` class and keeps this common benchmark source compilable on Spark 3.x. * + * Physical planning happens in [[prepareQuery]]. The SparkPlan it returns is the plan + * [[assertRoute]] inspects and the plan the timer executes. Do not use + * `DatasetToBenchmark.noop()` here: on Spark 4.x it calls `df.write.format("noop").save()`, which + * builds a fresh noop-write `QueryExecution` from the logical plan. That would mix write planning + * into the timed region and disconnect route validation from the action being measured. + * * To run this benchmark: * {{{ * SPARK_GENERATE_BENCHMARK_FILES=1 make \ @@ -226,16 +232,8 @@ object CometMapSortBenchmark extends CometBenchmarkBase { dispatch = false, shuffle = false, mapSortProjection().queryExecution.logical) - assertRoute( - shape, - "projection", - dispatch = true, - stripAQEPlan(dispatcher.queryExecution.executedPlan)) - assertRoute( - shape, - "projection", - dispatch = false, - stripAQEPlan(fallback.queryExecution.executedPlan)) + assertRoute(shape, "projection", dispatch = true, dispatcher.plan) + assertRoute(shape, "projection", dispatch = false, fallback.plan) runBenchmark(s"MapSort projection -- ${shape.description}") { val benchmark = new Benchmark( s"MapSort projection -- ${shape.description}", @@ -260,16 +258,8 @@ object CometMapSortBenchmark extends CometBenchmarkBase { dispatch = false, shuffle = true, shuffleQuery().queryExecution.logical) - assertRoute( - shape, - "shuffle", - dispatch = true, - stripAQEPlan(dispatcher.queryExecution.executedPlan)) - assertRoute( - shape, - "shuffle", - dispatch = false, - stripAQEPlan(fallback.queryExecution.executedPlan)) + assertRoute(shape, "shuffle", dispatch = true, dispatcher.plan) + assertRoute(shape, "shuffle", dispatch = false, fallback.plan) runBenchmark(s"MapSort native shuffle -- ${shape.description}") { val benchmark = new Benchmark( s"MapSort native shuffle -- ${shape.description}", @@ -471,41 +461,166 @@ object CometMapSortBenchmark extends CometBenchmarkBase { private def oneLine(value: String): String = value.split("\\n").iterator.map(_.trim).mkString(" | ") - /** Materializes physical planning outside Benchmark's timed closure. */ + /** + * A physical plan that has already been planned and route-validated, plus a consumer that + * executes that exact `SparkPlan` (or a fresh shuffle copy of it) without building another + * `QueryExecution`. + */ + private case class PreparedQuery(plan: SparkPlan, consume: SparkPlan => Unit) + + /** + * Forces physical planning outside the benchmark timer and returns that SparkPlan. The timed + * path must consume this plan (see [[executePreparedQuery]]) rather than calling + * `DatasetToBenchmark.noop()`, which on Spark 4.x creates a separate noop-write + * `QueryExecution`. + */ private def prepareQuery( shape: Shape, dispatch: Boolean, shuffle: Boolean, - logicalPlan: LogicalPlan): DataFrame = { - var prepared: DataFrame = null + logicalPlan: LogicalPlan): PreparedQuery = { + var prepared: PreparedQuery = null withSQLConf(configs(shape, dispatch, shuffle): _*) { - prepared = dataFrameOfRows(logicalPlan) - prepared.queryExecution.executedPlan + val df = dataFrameOfRows(logicalPlan) + val plan = stripAQEPlan(df.queryExecution.executedPlan) + val queryExecution = df.queryExecution + prepared = PreparedQuery( + plan, + consume = runnable => { + // Metrics / job grouping only. The body executes `runnable`, not a newly planned query. + SQLExecution.withNewExecutionId(queryExecution, Some("CometMapSortBenchmark")) { + consumePlan(runnable) + } + }) } prepared } - private def runPreparedQuery( + /** + * Executes and fully consumes an already-prepared physical plan. Does not construct a + * `QueryExecution`. + */ + private def executePreparedQuery( shape: Shape, dispatch: Boolean, shuffle: Boolean, - df: DataFrame): Unit = + prepared: PreparedQuery, + plan: SparkPlan): Unit = withSQLConf(configs(shape, dispatch, shuffle): _*) { - df.noop() + prepared.consume(plan) + } + + /** + * Consumes every partition of `plan` on the executors. Columnar plans stay columnar so the + * timer does not add a conversion that is absent from the validated tree. Row plans iterate + * every InternalRow. Closest in-repo precedent: `CometDatetimeExpressionBenchmark` drains + * `queryExecution.toRdd` with `foreachPartition` so Spark cannot skip execution. + */ + private def consumePlan(plan: SparkPlan): Unit = { + if (plan.supportsColumnar) { + plan.executeColumnar().foreachPartition { batches => + var rows = 0L + while (batches.hasNext) { + rows += batches.next().numRows() + } + if (rows < 0L) { + throw new IllegalStateException("columnar consume underflow") + } + } + } else { + plan.execute().foreachPartition { rows => + var n = 0L + while (rows.hasNext) { + rows.next() + n += 1L + } + if (n < 0L) { + throw new IllegalStateException("row consume underflow") + } + } + } + } + + /** + * `ShuffleExchangeExec` and `CometShuffleExchangeExec` cache their shuffle RDD / + * `ShuffleDependency`. Executing one instance twice would reuse shuffle files. Copying only + * those exchange nodes (children shared) yields a new shuffle id without repeating physical + * planning. + * + * `transformUp` / `withNewChildren` cannot be used here: both drop a `copy()` that `fastEquals` + * the original, and `CometShuffleExchangeExec.equals` compares partitioning and children, so + * the parent would keep the cached instance. Rebuild parents with `makeCopy` instead. A full + * `SparkPlan.clone()` is avoided because Comet native nodes can drop `@transient` scan state in + * `makeCopy`. + */ + private def withFreshShuffle(plan: SparkPlan): SparkPlan = { + def replaceExchanges(p: SparkPlan): SparkPlan = p match { + case comet: CometShuffleExchangeExec => comet.copy() + case sparkShuffle: ShuffleExchangeExec => sparkShuffle.copy() + case other => + val oldChildren = other.children + val newChildren = oldChildren.map(replaceExchanges) + if (newChildren.corresponds(oldChildren)(_ eq _)) { + other + } else { + replaceChildrenByIdentity(other, oldChildren, newChildren) + } } + val copied = replaceExchanges(plan) + val originalExchanges = plan.collect { + case s: CometShuffleExchangeExec => s + case s: ShuffleExchangeExec => s + } + val copiedExchanges = copied.collect { + case s: CometShuffleExchangeExec => s + case s: ShuffleExchangeExec => s + } + assert( + originalExchanges.nonEmpty && + originalExchanges.length == copiedExchanges.length && + originalExchanges.zip(copiedExchanges).forall { case (left, right) => !(left eq right) }, + s"expected a fresh shuffle exchange instance:\n${plan.treeString}") + copied + } + + /** Replaces children by reference equality so an equal-but-new exchange is not discarded. */ + private def replaceChildrenByIdentity( + plan: SparkPlan, + oldChildren: Seq[SparkPlan], + newChildren: Seq[SparkPlan]): SparkPlan = { + var idx = 0 + val newArgs = plan.productIterator.map { + case child: SparkPlan if idx < oldChildren.length && (child eq oldChildren(idx)) => + val replacement = newChildren(idx) + idx += 1 + replacement + case other => other.asInstanceOf[AnyRef] + }.toArray + assert( + idx == oldChildren.length, + s"could not rebuild children for ${plan.nodeName}:\n${plan.treeString}") + plan.makeCopy(newArgs).asInstanceOf[SparkPlan] + } + private def addMatchedCases( benchmark: Benchmark, shape: Shape, shuffle: Boolean, - fallback: DataFrame, - dispatcher: DataFrame): Unit = { - def addFallback(): Unit = benchmark.addCase(FallbackCaseName) { _ => - runPreparedQuery(shape, dispatch = false, shuffle = shuffle, df = fallback) - } - def addDispatcher(): Unit = benchmark.addCase(DispatcherCaseName) { _ => - runPreparedQuery(shape, dispatch = true, shuffle = shuffle, df = dispatcher) + fallback: PreparedQuery, + dispatcher: PreparedQuery): Unit = { + def addArm(name: String, dispatch: Boolean, prepared: PreparedQuery): Unit = { + // Copy shuffle exchanges before startTiming so query planning stays outside the + // measured region while every sample still performs a fresh shuffle. + benchmark.addTimerCase(name) { timer => + val runnable = if (shuffle) withFreshShuffle(prepared.plan) else prepared.plan + timer.startTiming() + executePreparedQuery(shape, dispatch, shuffle, prepared, runnable) + timer.stopTiming() + } } + def addFallback(): Unit = addArm(FallbackCaseName, dispatch = false, fallback) + def addDispatcher(): Unit = addArm(DispatcherCaseName, dispatch = true, dispatcher) CaseOrder match { case "fallback-first" => @@ -558,20 +673,27 @@ object CometMapSortBenchmark extends CometBenchmarkBase { val logicalPlan = if (shuffle) shuffleQuery().queryExecution.logical else mapSortProjection().queryExecution.logical - val df = prepareQuery(shape, dispatch, shuffle, logicalPlan) + val prepared = prepareQuery(shape, dispatch, shuffle, logicalPlan) + // Inspect the plan that will be timed. Do not execute it here: that would warm the + // codegen dispatcher before the first-action sample. + assertRoute(shape, workload, dispatch, prepared.plan) CometScalaUDFCodegen.resetStats() val elapsed = - timeMillis(runPreparedQuery(shape, dispatch = dispatch, shuffle = shuffle, df = df)) - val plan = stripAQEPlan(df.queryExecution.executedPlan) + timeMillis( + executePreparedQuery( + shape, + dispatch = dispatch, + shuffle = shuffle, + prepared = prepared, + plan = prepared.plan)) val stats = CometScalaUDFCodegen.stats() - assertRoute(shape, workload, dispatch, plan) val repetition = sys.props.getOrElse(RepetitionProperty, "") emit( f"MAPSORT_FIRST_ACTION shape=${shape.name} route=$route workload=$workload " + f"repetition=$repetition rows=$FirstActionRows elapsed_ms=$elapsed%.1f " + s"dispatcher_compile_count=${stats.compileCount} " + s"dispatcher_cache_hit_count=${stats.cacheHitCount}") - emit(s" executed plan: ${oneLine(plan.treeString)}") + emit(s" executed plan: ${oneLine(prepared.plan.treeString)}") emit( "This is first-action latency from one process. Dispatcher counters are routing/cache " + "observations, not a measurement of Java compiler time.") From 0f0bcc124ceca7eef8dc6fb6ef2511c25c9bc939 Mon Sep 17 00:00:00 2001 From: Hung Date: Mon, 14 Sep 2026 23:56:51 +0800 Subject: [PATCH 4/8] test: strengthen MapSort dispatcher correctness coverage --- .../comet/CometMapExpressionSuite.scala | 140 ++++++++++++++-- .../comet/exec/CometNativeShuffleSuite.scala | 157 +++++++++++++++--- .../sql/benchmark/CometMapSortBenchmark.scala | 51 ++++-- 3 files changed, 306 insertions(+), 42 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index c38b8a56d69..36b3db009f4 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -30,6 +30,7 @@ import org.apache.spark.sql.types.BinaryType import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions} +import org.apache.comet.udf.codegen.CometScalaUDFCodegen class CometMapExpressionSuite extends CometTestBase with CometCodegenAssertions { @@ -325,40 +326,141 @@ class CometMapExpressionSuite extends CometTestBase with CometCodegenAssertions } } - test("mapsort with complex keys falls back when codegen dispatcher is disabled") { + test("mapsort expression disable restores fallback while codegen dispatcher stays enabled") { assume(isSpark40Plus, "Spark 4.0 inserts MapSort for group-by on map keys") withTable("t_map_sort_dispatch_disabled") { sql("CREATE TABLE t_map_sort_dispatch_disabled (m MAP, INT>) USING parquet") sql("""INSERT INTO t_map_sort_dispatch_disabled VALUES |(map(array(2, 1), 20, array(1, 2), 10)), |(map(array(1, 2), 10, array(2, 1), 20))""".stripMargin) - withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") { + withSQLConf( + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", + CometConf.getExprEnabledConfigKey("MapSort") -> "false") { val df = sql("SELECT m, count(*) FROM t_map_sort_dispatch_disabled GROUP BY m") assertMapSortInPlan(df) assertCodegenDidNotRun { checkSparkAnswerAndFallbackReason( df, - CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key + "=false") + "Expression support is disabled. Set " + + s"${CometConf.getExprEnabledConfigKey("MapSort")}=true to enable it.") } } } } - // Spark rejects collated strings as map keys (`UNSUPPORTED_FEATURE.COLLATIONS_IN_MAP_KEYS`), so - // `MapSort` never sees that shape. `supportedScalarSortElementType` still excludes them, and the - // same `Unsupported` → dispatcher path is covered by the array/struct cases above. + test("mapsort routes collated-string keys through codegen dispatcher") { + assume(isSpark40Plus, "collated map keys and MapSort require Spark 4.0+") + withSQLConf( + "spark.sql.collation.allowInMapKeys" -> "true", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + withTable("t_map_sort_collated_key") { + sql( + "CREATE TABLE t_map_sort_collated_key " + + "(k1 STRING, v1 INT, k2 STRING, v2 INT) USING parquet") + sql("""INSERT INTO t_map_sort_collated_key VALUES + |('b', 20, 'A', 10), + |('a', 10, 'B', 20), + |('c', 30, 'd', 40)""".stripMargin) + val query = + """SELECT m, count(*) FROM ( + | SELECT map(CAST(k1 AS STRING COLLATE UTF8_LCASE), v1, + | CAST(k2 AS STRING COLLATE UTF8_LCASE), v2) AS m + | FROM t_map_sort_collated_key) + |GROUP BY m""".stripMargin + val df = sql(query) + + assertMapSortInPlan(df) + CometScalaUDFCodegen.resetStats() + val cometRows = df.collect() + val cometPlan = df.queryExecution.executedPlan + val dispatcherStats = CometScalaUDFCodegen.stats() + assert( + dispatcherStats.totalLookups >= 1, + s"expected codegen dispatcher activity, got $dispatcherStats; " + + s"fallback reasons: ${new ExtendedExplainInfo().getFallbackReasons(cometPlan)}\n" + + cometPlan) + + // UTF8_LCASE considers the first two maps equal but Spark and Comet may retain different + // byte-level representatives for the grouped key ("A" versus "a"). Compare the collected + // answers after canonicalizing keys in the test, leaving the executed SQL plan untouched. + def canonicalize(rows: Array[org.apache.spark.sql.Row]) = + rows + .map { row => + val entries = + if (row.isNullAt(0)) { + None + } else { + Some( + row + .getMap[String, Int](0) + .toSeq + .map { case (key, value) => + key.toLowerCase(java.util.Locale.ROOT) -> value + } + .sortBy(_._1)) + } + entries -> row.getLong(1) + } + .sortBy(_.toString) + .toSeq + + var sparkRows: Array[org.apache.spark.sql.Row] = Array.empty + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkRows = sql(query).collect() + } + assert(canonicalize(cometRows) === canonicalize(sparkRows)) + + val explain = new ExtendedExplainInfo() + assert( + explain.getCodegenDispatchExpressions(cometPlan).contains("mapsort"), + s"expected collated-key mapsort on codegen dispatch path:\n$cometPlan") + assert( + !explain.getNativeExpressions(cometPlan).contains("mapsort"), + s"collated-key mapsort must not use native map_sort:\n$cometPlan") + + withSQLConf(CometConf.getExprEnabledConfigKey("MapSort") -> "false") { + val fallback = sql(query) + assertMapSortInPlan(fallback) + val fallbackRows = assertCodegenDidNotRun(fallback.collect()) + val fallbackPlan = fallback.queryExecution.executedPlan + assert(canonicalize(fallbackRows) === canonicalize(sparkRows)) + val expectedReason = + "Expression support is disabled. Set " + + s"${CometConf.getExprEnabledConfigKey("MapSort")}=true to enable it." + assert( + new ExtendedExplainInfo().getFallbackReasons(fallbackPlan).contains(expectedReason), + s"expected MapSort-specific fallback reason `$expectedReason` in:\n$fallbackPlan") + } + } + } + } test("mapsort routes strict floating-point keys through codegen dispatcher") { assume(isSpark40Plus, "Spark 4.0 inserts MapSort for group-by on map keys") - withSQLConf(CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") { + withSQLConf( + CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", + "spark.sql.legacy.disableMapKeyNormalization" -> "true") { withTable("t_map_sort_fp_key") { - sql("CREATE TABLE t_map_sort_fp_key (m MAP) USING parquet") + sql("CREATE TABLE t_map_sort_fp_key (id INT, m MAP) USING parquet") sql("""INSERT INTO t_map_sort_fp_key VALUES - |(map(CAST('NaN' AS DOUBLE), 1, CAST('-0.0' AS DOUBLE), 2, 1.0, 3)), - |(map(1.0, 3, CAST('-0.0' AS DOUBLE), 2, CAST('NaN' AS DOUBLE), 1)), - |(map(0.0, 4)), - |(NULL)""".stripMargin) + |(1, map(CAST('0.0' AS DOUBLE), 10, CAST('-0.0' AS DOUBLE), 20, + | CAST('NaN' AS DOUBLE), 30)), + |(2, map(CAST('-0.0' AS DOUBLE), 20, CAST('0.0' AS DOUBLE), 10, + | CAST('NaN' AS DOUBLE), 30)), + |(3, map(CAST('NaN' AS DOUBLE), 40, 1.0, 50)), + |(4, NULL)""".stripMargin) + + val storedKeys = sql( + "SELECT k FROM t_map_sort_fp_key " + + "LATERAL VIEW explode(map_keys(m)) e AS k WHERE id = 1").collect().map(_.getDouble(0)) + val storedBits = storedKeys.map(java.lang.Double.doubleToRawLongBits) + assert( + storedBits.take(2).sameElements(Array(0L, Long.MinValue)), + s"expected stored +0.0 then -0.0 raw bits, got ${storedBits.toSeq}") + assert(storedKeys.exists(java.lang.Double.isNaN), "strict fixture must retain a NaN key") + val df = sql("SELECT m, count(*) FROM t_map_sort_fp_key GROUP BY m") assertMapSortInPlan(df) @@ -369,6 +471,20 @@ class CometMapExpressionSuite extends CometTestBase with CometCodegenAssertions assert( dispatched.contains("mapsort"), s"expected mapsort on codegen dispatch path, got $dispatched in:\n$cometPlan") + assert( + !new ExtendedExplainInfo().getNativeExpressions(cometPlan).contains("mapsort"), + s"strict floating-point mapsort must not use native map_sort:\n$cometPlan") + + withSQLConf(CometConf.getExprEnabledConfigKey("MapSort") -> "false") { + val fallback = sql("SELECT m, count(*) FROM t_map_sort_fp_key GROUP BY m") + assertMapSortInPlan(fallback) + assertCodegenDidNotRun { + checkSparkAnswerAndFallbackReason( + fallback, + "Expression support is disabled. Set " + + s"${CometConf.getExprEnabledConfigKey("MapSort")}=true to enable it.") + } + } } } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index b0be7856038..c3b4654fe8d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -49,7 +49,7 @@ import org.apache.spark.sql.functions.{col, count, sum} import org.apache.spark.sql.types.{ArrayType, DataType, LongType, MapType, StructField, StructType} import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} -import org.apache.comet.{CometCodegenAssertions, CometConf, CometExecIterator, CometExplainInfo, CometShuffleBlockIterator, CometShuffleSizeLimitException, Native} +import org.apache.comet.{CometCodegenAssertions, CometConf, CometExecIterator, CometExplainInfo, CometShuffleBlockIterator, CometShuffleSizeLimitException, ExtendedExplainInfo, Native} import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus import org.apache.comet.serde.{OperatorOuterClass, PartitioningOuterClass} import org.apache.comet.shuffle.ShufflePartitionPusher @@ -437,7 +437,7 @@ class CometNativeShuffleSuite withTempDir { dir => val path = new Path(dir.toURI.toString, "test.parquet") makeParquetFileAllPrimitiveTypes(path, dictionaryEnabled = dictionaryEnabled, 1000) - var allTypes: Seq[Int] = (1 to 20) + val allTypes: Seq[Int] = (1 to 20) allTypes.map(i => s"_$i").foreach { c => withSQLConf("parquet.enable.dictionary" -> dictionaryEnabled.toString) { readParquetFile(path.toString) { df => @@ -958,8 +958,6 @@ class CometNativeShuffleSuite "SELECT _1, spark_partition_id() AS pid FROM (" + s"SELECT /*+ REPARTITION(10, $keys) */ * FROM tbl)" val cometRows = sql(query).collect().map(r => (r.getInt(0), r.getInt(1))).sorted - // `withSQLConf` returns Unit, so capture the Spark-side rows via a var rather than - // relying on the block's value. // `SQLHelper.withSQLConf` returns T on Spark 4.x but Unit on Spark 3.x, so capture the // Spark-side rows via a var to keep this compiling on both. var sparkRows: Array[(Int, Int)] = Array.empty @@ -976,6 +974,81 @@ class CometNativeShuffleSuite } } + if (isSpark40Plus) { + withTable("complex_map_partition_keys") { + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sql("""CREATE TABLE complex_map_partition_keys USING parquet AS + |SELECT CAST(id AS INT) AS id, + | CASE WHEN pmod(id, 2) = 0 THEN + | map(array(CAST(pmod(id, 13) AS INT), 1), CAST(pmod(id, 17) AS INT), + | array(CAST(pmod(id, 7) AS INT), 2), CAST(pmod(id, 19) AS INT)) + | ELSE + | map(array(CAST(pmod(id, 7) AS INT), 2), CAST(pmod(id, 19) AS INT), + | array(CAST(pmod(id, 13) AS INT), 1), CAST(pmod(id, 17) AS INT)) + | END AS array_map, + | CASE WHEN pmod(id, 2) = 0 THEN + | map(named_struct('a', CAST(pmod(id, 13) AS INT), 'b', 'x'), + | CAST(pmod(id, 17) AS INT), + | named_struct('a', CAST(pmod(id, 7) AS INT), 'b', 'y'), + | CAST(pmod(id, 19) AS INT)) + | ELSE + | map(named_struct('a', CAST(pmod(id, 7) AS INT), 'b', 'y'), + | CAST(pmod(id, 19) AS INT), + | named_struct('a', CAST(pmod(id, 13) AS INT), 'b', 'x'), + | CAST(pmod(id, 17) AS INT)) + | END AS struct_map + |FROM range(200)""".stripMargin) + } + + assert( + CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_NESTED_ENABLED + .get(spark.sessionState.conf), + "complex map-key parity requires nested hash partitioning to be enabled") + + Seq("array_map", "struct_map").foreach { key => + val query = + "SELECT id, spark_partition_id() AS pid FROM (" + + s"SELECT /*+ REPARTITION(10, $key) */ id, $key " + + "FROM complex_map_partition_keys)" + val cometDf = sql(query) + val cometRows = assertCodegenRan { + cometDf.collect().map(r => (r.getInt(0), r.getInt(1))).sorted + } + val cometPlan = stripAQEPlan(cometDf.queryExecution.executedPlan) + val cometExchanges = collect(cometPlan) { case e: CometShuffleExchangeExec => e } + assert( + cometExchanges.size == 1, + s"$key expected one CometShuffleExchangeExec:\n$cometPlan") + assert( + new ExtendedExplainInfo() + .getCodegenDispatchExpressions(cometPlan) + .contains("mapsort"), + s"$key expected MapSort on the JVM dispatcher path:\n$cometPlan") + + var sparkRows: Array[(Int, Int)] = Array.empty + var sparkPartitions = -1 + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + val sparkDf = sql(query) + sparkRows = sparkDf.collect().map(r => (r.getInt(0), r.getInt(1))).sorted + val sparkPlan = stripAQEPlan(sparkDf.queryExecution.executedPlan) + val sparkExchanges = collect(sparkPlan) { case e: ShuffleExchangeExec => e } + assert( + sparkExchanges.size == 1, + s"$key expected one Spark ShuffleExchangeExec:\n$sparkPlan") + sparkPartitions = sparkExchanges.head.outputPartitioning.numPartitions + } + + assert(sparkRows.nonEmpty, s"Spark produced no rows for $key") + assert( + cometExchanges.head.outputPartitioning.numPartitions == sparkPartitions, + s"$key output partition count differs from Spark") + assert( + cometRows === sparkRows, + s"partition assignment differs from Spark for map key $key") + } + } + } + // Same check for a deeply nested key, where Spark rewrites the partitioning expression into a // transform(...) containing a nested mapsort(...). if (isSpark40Plus) { @@ -1034,23 +1107,69 @@ class CometNativeShuffleSuite // and assert both rows land in the same partition. assume(isSpark40Plus, "map shuffle keys are only normalized with mapsort on Spark 4.0+") withSQLConf(CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { - val query = """ - SELECT id, spark_partition_id() AS pid FROM ( - SELECT /*+ REPARTITION(10, m) */ * FROM VALUES - (1, map('a', 1, 'b', 2)), - (2, map('b', 2, 'a', 1)), - (3, map('a', 1, 'b', 2, 'c', 3)), - (4, map('c', 3, 'b', 2, 'a', 1)) AS t(id, m))""" - val rows = sql(query).collect().map(r => (r.getInt(0), r.getInt(1))).toMap - assert(rows(1) == rows(2), "equal maps in different entry order must share a partition") - assert(rows(3) == rows(4), "equal maps in different entry order must share a partition") + val cases = Seq( + ( + "scalar", + """(1, map('a', 1, 'b', 2)), + |(2, map('b', 2, 'a', 1)), + |(3, map('a', 1, 'b', 2, 'c', 3)), + |(4, map('c', 3, 'b', 2, 'a', 1))""".stripMargin, + false), + ( + "array", + """(1, map(array(1, 2), 10, array(2, 1), 20)), + |(2, map(array(2, 1), 20, array(1, 2), 10)), + |(3, map(array(1), 10, array(2), 20, array(3), 30)), + |(4, map(array(3), 30, array(2), 20, array(1), 10))""".stripMargin, + true), + ( + "struct", + """(1, map(named_struct('a', 1, 'b', 'x'), 10, + | named_struct('a', 2, 'b', 'y'), 20)), + |(2, map(named_struct('a', 2, 'b', 'y'), 20, + | named_struct('a', 1, 'b', 'x'), 10)), + |(3, map(named_struct('a', 1, 'b', 'x'), 10, + | named_struct('a', 2, 'b', 'y'), 20, + | named_struct('a', 3, 'b', 'z'), 30)), + |(4, map(named_struct('a', 3, 'b', 'z'), 30, + | named_struct('a', 2, 'b', 'y'), 20, + | named_struct('a', 1, 'b', 'x'), 10))""".stripMargin, + true)) + + cases.foreach { case (name, values, expectDispatch) => + val query = s"""SELECT id, spark_partition_id() AS pid FROM ( + | SELECT /*+ REPARTITION(10, m) */ * FROM VALUES + | $values AS t(id, m))""".stripMargin + val df = sql(query) + val rows = { + def collectRows(): Map[Int, Int] = + df.collect().map(r => (r.getInt(0), r.getInt(1))).toMap + if (expectDispatch) assertCodegenRan(collectRows()) else collectRows() + } + assert( + rows(1) == rows(2), + s"equal $name-keyed maps in different entry order must share a partition") + assert( + rows(3) == rows(4), + s"equal $name-keyed maps in different entry order must share a partition") - var sparkRows: Map[Int, Int] = Map.empty - withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - sparkRows = sql(query).collect().map(r => (r.getInt(0), r.getInt(1))).toMap + val cometPlan = stripAQEPlan(df.queryExecution.executedPlan) + assert( + collect(cometPlan) { case e: CometShuffleExchangeExec => e }.nonEmpty, + s"$name-keyed map did not retain native shuffle:\n$cometPlan") + val dispatched = + new ExtendedExplainInfo().getCodegenDispatchExpressions(cometPlan).contains("mapsort") + assert( + dispatched == expectDispatch, + s"unexpected mapsort dispatch route for $name-keyed map:\n$cometPlan") + + var sparkRows: Map[Int, Int] = Map.empty + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkRows = sql(query).collect().map(r => (r.getInt(0), r.getInt(1))).toMap + } + assert(sparkRows.nonEmpty) + assert(rows == sparkRows, s"$name-keyed map partition assignment differs from Spark") } - assert(sparkRows.nonEmpty) - assert(rows == sparkRows, "map key partition assignment differs from Spark") } } } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala index 591f8596230..e13704efc7f 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala @@ -38,14 +38,15 @@ import org.apache.comet.udf.codegen.CometScalaUDFCodegen * Matched benchmark for the two routes Spark 4.x can take for `MapSort` shapes that Comet cannot * sort natively: * - * - with the JVM codegen dispatcher disabled, the enclosing projection or shuffle falls back to - * Spark; and - * - with it enabled, Spark's `MapSort.doGenCode` executes inside the Comet pipeline. + * - with `spark.comet.expression.MapSort.enabled=false`, the enclosing projection or shuffle + * falls back to Spark; and + * - with MapSort enabled, Spark's `MapSort.doGenCode` executes inside the Comet pipeline. * * Every pair reads the same Parquet data and differs only in - * `spark.comet.exec.scalaUDF.codegen.enabled`. Array and struct cases vary map size independently - * from nested-key width; strict floating-point cases include NaN and signed zero. Input maps are - * written in reverse key order and one row in 64 has a NULL map. + * `spark.comet.expression.MapSort.enabled`; the global codegen dispatcher remains enabled in both + * arms. Array and struct cases vary map size independently from nested-key width. Strict + * floating-point cases include NaN and both signed zeros in the same map. Input maps are written + * in reverse key order and one row in 64 has a NULL map. * * Spark 4.0 and 4.1 only insert `MapSort` for grouping and repartition expressions; * `try_element_at` itself does not insert one. To measure a projection without also timing an @@ -159,12 +160,13 @@ object CometMapSortBenchmark extends CometBenchmarkBase { s"""CASE | WHEN $entry = 0 THEN CAST('NaN' AS DOUBLE) | WHEN $entry = 1 THEN CAST('-0.0' AS DOUBLE) + | WHEN $entry = 2 THEN CAST('0.0' AS DOUBLE) | ELSE CAST(id * 128 + CAST($entry AS BIGINT) + 1 AS DOUBLE) |END""".stripMargin.replace('\n', ' ') } private case class Shape(name: String, family: KeyFamily, mapSize: Int, keyWidth: Int) { - require(mapSize >= 2, "mapSize must leave room for NaN and -0.0") + require(mapSize >= 3, "mapSize must leave room for NaN, -0.0, and +0.0") require(keyWidth > 0, "keyWidth must be positive") def description: String = @@ -639,9 +641,12 @@ object CometMapSortBenchmark extends CometBenchmarkBase { Seq( CometConf.COMET_ENABLED.key -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true", - CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> dispatch.toString, + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", + CometConf.getExprEnabledConfigKey("MapSort") -> dispatch.toString, CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> (shape.family == StrictDoubleKey).toString, + "spark.sql.legacy.disableMapKeyNormalization" -> + (shape.family == StrictDoubleKey).toString, CometConf.getExprAllowIncompatConfigKey("MapSort") -> "false", CometConf.COMET_SHUFFLE_ENABLED.key -> shuffle.toString, CometConf.COMET_SHUFFLE_MODE.key -> "native", @@ -712,12 +717,35 @@ object CometMapSortBenchmark extends CometBenchmarkBase { spark .range(0L, rows.toLong, 1L, InputPartitions) .createOrReplaceTempView(tbl) - prepareTable(dir, spark.sql(corpusQuery(shape))) - f + withSQLConf( + "spark.sql.legacy.disableMapKeyNormalization" -> + (shape.family == StrictDoubleKey).toString) { + prepareTable(dir, spark.sql(corpusQuery(shape))) + assertStrictDoubleCorpus(shape) + f + } } } } + private def assertStrictDoubleCorpus(shape: Shape): Unit = { + if (shape.family == StrictDoubleKey) { + val keys = spark + .sql( + "SELECT key FROM parquetV1Table " + + "LATERAL VIEW explode(map_keys(m)) e AS key WHERE id = 1") + .collect() + .map(_.getDouble(0)) + val zeroBits = keys + .filter(_ == 0.0d) + .map(java.lang.Double.doubleToRawLongBits) + assert( + zeroBits.sameElements(Array(0L, Long.MinValue)), + s"${shape.description}: expected stored +0.0 then -0.0, got ${zeroBits.toSeq}") + assert(keys.exists(java.lang.Double.isNaN), s"${shape.description}: missing stored NaN key") + } + } + private def corpusQuery(shape: Shape): String = { val entrySequence = s"sequence(${shape.mapSize - 1}, 0, -1)" val keys = @@ -755,7 +783,8 @@ object CometMapSortBenchmark extends CometBenchmarkBase { emit(s"Correctness/routing prefix rows: $VerificationRows; NULL map density: 1/$NullMapEvery") emit(s"Mode: $Mode; steady-state case order: $CaseOrder") emit("Shuffle mode: native; AQE: disabled; nested hash partitioning: enabled") - emit(s"Only matched-arm difference: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}") + emit(s"Global dispatcher enabled: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=true") + emit(s"Only matched-arm difference: ${CometConf.getExprEnabledConfigKey("MapSort")}") emit(s"Pinned: ${CometConf.getExprAllowIncompatConfigKey("MapSort")}=false") emit("Projection note: Spark 4.0/4.1 do not insert MapSort for try_element_at. This suite") emit(" executes the optimizer-inserted grouping MapSort Project alone; no aggregate is run.") From 5275b4397df5fa6b78c481b7b884b2d5601e4486 Mon Sep 17 00:00:00 2001 From: Hung Date: Mon, 14 Sep 2026 23:57:05 +0800 Subject: [PATCH 5/8] docs: clarify MapSort JVM routing and rollback --- .../expressions/_category_template/map.md | 16 +++++++++++----- .../org/apache/comet/serde/CometMapSort.scala | 12 +++++++++--- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/expressions/_category_template/map.md b/docs/source/user-guide/latest/compatibility/expressions/_category_template/map.md index 1676e4542da..3bbfa7ff5f0 100644 --- a/docs/source/user-guide/latest/compatibility/expressions/_category_template/map.md +++ b/docs/source/user-guide/latest/compatibility/expressions/_category_template/map.md @@ -21,11 +21,11 @@ under the License. ## MapSort (Spark 4.0+) -Spark 4.0 inserts `MapSort` to normalize map values when they appear in shuffle hash partitioning -keys, in `try_element_at`, and in other contexts where map ordering must be deterministic. Comet -runs `MapSort` natively for supported scalar key types. Other orderable key types use Spark's own -generated code through the JVM codegen dispatcher, so the enclosing operator stays in the Comet -pipeline. +Spark 4.0 inserts `MapSort` to normalize map values when they appear in grouping expressions or +shuffle hash partitioning keys. Comet runs `MapSort` natively for supported scalar key types. +Other dispatcher-eligible orderable key types use Spark's own generated JVM code through the +codegen dispatcher, so the enclosing operator can stay in the Comet pipeline. This dispatcher +route is not a native `MapSort` implementation. When `spark.comet.exec.strictFloatingPoint=true`, maps whose keys contain `Float` or `Double` also use the dispatcher (consistent with `SortOrder` and `SortArray`). Arrow's sort uses IEEE total @@ -33,5 +33,11 @@ ordering for floating-point, which differs from Spark's `Double.compare` semanti `-0.0`. If the dispatcher is disabled or cannot handle an expression, Comet safely falls back to Spark. +Set `spark.comet.expression.MapSort.enabled=false` to restore the previous behavior, where a +`MapSort` without a native implementation causes its enclosing projection or shuffle to fall back +to Spark. This expression-specific setting leaves the codegen dispatcher available to unrelated +expressions. Retaining an enclosing operator in Comet is a functional routing benefit; by itself, +it does not guarantee higher throughput. + diff --git a/spark/src/main/spark-4.x/org/apache/comet/serde/CometMapSort.scala b/spark/src/main/spark-4.x/org/apache/comet/serde/CometMapSort.scala index f238ccd4068..1a48d099196 100644 --- a/spark/src/main/spark-4.x/org/apache/comet/serde/CometMapSort.scala +++ b/spark/src/main/spark-4.x/org/apache/comet/serde/CometMapSort.scala @@ -25,7 +25,11 @@ import org.apache.spark.sql.types.MapType import org.apache.comet.CometConf import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, scalarFunctionExprToProtoWithReturnType, supportedScalarSortElementType} -// Key types unsupported by the native kernel can still run Spark's generated code in-pipeline. +// Key types without a native implementation can still run Spark's generated code in-pipeline. +// Spark rejects collated-string map keys by default, but they reach MapSort when +// spark.sql.collation.allowInMapKeys=true and the map is built from dispatcher-supported inputs; +// those expressions take the same Unsupported -> dispatcher route. A scan carrying a collated +// map schema may still be rejected independently by the scan's schema support checks. object CometMapSort extends CometExpressionSerde[MapSort] with CodegenDispatchFallback { override def getIncompatibleReasons(): Seq[String] = @@ -34,12 +38,14 @@ object CometMapSort extends CometExpressionSerde[MapSort] with CodegenDispatchFa s"`${CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key}=true`.") override def getUnsupportedReasons(): Seq[String] = - Seq("MapSort is unsupported for non-scalar key types (struct, array, map, etc.).") + Seq( + "MapSort with an orderable key type outside native scalar coverage, including array, " + + "struct, interval, and non-default-collated string keys, has no native implementation.") override def getSupportLevel(expr: MapSort): SupportLevel = { val keyType = expr.dataType.asInstanceOf[MapType].keyType if (!supportedScalarSortElementType(keyType)) { - Unsupported(Some(s"MapSort on map with key type $keyType is not supported")) + Unsupported(Some(s"MapSort with key type $keyType has no native implementation")) } else { SupportLevel .strictFloatingPointReason(keyType, "MapSort on floating-point key") From db2a64d89a447064b2aea353db5fac52bd3b543f Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Mon, 21 Sep 2026 17:13:32 +0800 Subject: [PATCH 6/8] feat: make the MapSort codegen dispatcher opt-in by default Matched microbenchmarks show the dispatcher slower than Spark fallback for most projection shapes and all shuffle shapes, so keep native scalar map_sort on and require MapSort.codegen.enabled=true to stay in the Comet pipeline. --- .../expressions/_category_template/map.md | 32 +-- .../scala/org/apache/comet/CometConf.scala | 15 ++ .../scala/org/apache/comet/GenerateDocs.scala | 49 ++++- .../comet/serde/CometExpressionSerde.scala | 26 ++- .../apache/comet/serde/QueryPlanSerde.scala | 13 +- .../shuffle/CometShuffleExchangeExec.scala | 10 +- .../org/apache/comet/serde/CometMapSort.scala | 19 +- .../comet/CometMapExpressionSuite.scala | 92 +++++--- .../org/apache/comet/GenerateDocsSuite.scala | 28 +++ .../exec/CometColumnarShuffleSuite.scala | 32 ++- .../comet/exec/CometNativeShuffleSuite.scala | 205 +++++++++++------- .../sql/benchmark/CometMapSortBenchmark.scala | 23 +- 12 files changed, 374 insertions(+), 170 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/expressions/_category_template/map.md b/docs/source/user-guide/latest/compatibility/expressions/_category_template/map.md index 3bbfa7ff5f0..c45efa04b7e 100644 --- a/docs/source/user-guide/latest/compatibility/expressions/_category_template/map.md +++ b/docs/source/user-guide/latest/compatibility/expressions/_category_template/map.md @@ -23,21 +23,23 @@ under the License. Spark 4.0 inserts `MapSort` to normalize map values when they appear in grouping expressions or shuffle hash partitioning keys. Comet runs `MapSort` natively for supported scalar key types. -Other dispatcher-eligible orderable key types use Spark's own generated JVM code through the -codegen dispatcher, so the enclosing operator can stay in the Comet pipeline. This dispatcher -route is not a native `MapSort` implementation. - -When `spark.comet.exec.strictFloatingPoint=true`, maps whose keys contain `Float` or `Double` also -use the dispatcher (consistent with `SortOrder` and `SortArray`). Arrow's sort uses IEEE total -ordering for floating-point, which differs from Spark's `Double.compare` semantics for `NaN` and -`-0.0`. If the dispatcher is disabled or cannot handle an expression, Comet safely falls back to -Spark. - -Set `spark.comet.expression.MapSort.enabled=false` to restore the previous behavior, where a -`MapSort` without a native implementation causes its enclosing projection or shuffle to fall back -to Spark. This expression-specific setting leaves the codegen dispatcher available to unrelated -expressions. Retaining an enclosing operator in Comet is a functional routing benefit; by itself, -it does not guarantee higher throughput. + +Other orderable key types have no native implementation. By default those cases fall the +enclosing projection or shuffle back to Spark. Set `spark.comet.expression.MapSort.codegen.enabled=true` +to run Spark's own generated JVM code through the codegen dispatcher instead, so the enclosing +operator can stay in the Comet pipeline. That dispatcher route is not a native `MapSort` +implementation, and a matched microbenchmark of these shapes is slower than the Spark fallback. + +When `spark.comet.exec.strictFloatingPoint=true`, maps whose keys contain `Float` or `Double` +follow the same default: the enclosing operator falls back to Spark. The codegen setting above +routes those keys through the dispatcher instead (consistent with `SortOrder` and `SortArray`). +Arrow's sort uses IEEE total ordering for floating-point, which differs from Spark's +`Double.compare` semantics for `NaN` and `-0.0`. If the dispatcher is disabled or cannot handle +an expression, Comet safely falls back to Spark. + +Set `spark.comet.expression.MapSort.enabled=false` to disable Comet `MapSort` entirely, including +the native scalar-key path. That expression-specific setting leaves the codegen dispatcher +available to unrelated expressions. diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index b4455f63c4c..e38bb8dc5ce 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -411,6 +411,21 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(true) + val COMET_EXPRESSION_MAPSORT_CODEGEN_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.expression.MapSort.codegen.enabled") + .category(CATEGORY_EXEC) + .doc( + "Whether to route `MapSort` cases that have no native implementation through Comet's " + + "JVM codegen dispatcher (Spark's own `doGenCode` inside the Comet pipeline). When " + + "enabled, array, struct, interval, and non-default-collated string map keys, and " + + "floating-point keys under `spark.comet.exec.strictFloatingPoint=true`, keep their " + + "enclosing projection or shuffle in Comet. Disabled by default because a matched " + + "microbenchmark of those shapes is slower than falling the enclosing operator back " + + "to Spark. Scalar map keys continue to use native `map_sort` regardless of this " + + "setting. This flag does not change codegen dispatch for any other expression.") + .booleanConf + .createWithDefault(false) + val COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.shuffle.native.partitioning.hash.enabled") .withAlternative("spark.comet.native.shuffle.partitioning.hash.enabled") diff --git a/spark/src/main/scala/org/apache/comet/GenerateDocs.scala b/spark/src/main/scala/org/apache/comet/GenerateDocs.scala index 4eb29a6d58d..900e1561dd4 100644 --- a/spark/src/main/scala/org/apache/comet/GenerateDocs.scala +++ b/spark/src/main/scala/org/apache/comet/GenerateDocs.scala @@ -62,7 +62,12 @@ object GenerateDocs { * retain the native opt-in behavior * @param codegenDispatchFallback * whether the serde mixes in `CodegenDispatchFallback`, meaning `unsupportedReasons` cases - * route through the JVM codegen dispatcher instead of falling back to Spark + * can route through the JVM codegen dispatcher instead of falling back to Spark + * @param codegenDispatchEnabledByDefault + * whether those dispatcher cases run by default. When false, they fall back to Spark until + * the user sets `codegenDispatchConfigKey` + * @param codegenDispatchConfigKey + * the config key that enables codegen dispatch when it is not the default */ private[comet] case class ExprNotes( name: String, @@ -72,7 +77,9 @@ object GenerateDocs { nativeOptIn: Boolean, nativeOptInConfigKey: String, conditionalNativeDefault: Boolean, - codegenDispatchFallback: Boolean) + codegenDispatchFallback: Boolean, + codegenDispatchEnabledByDefault: Boolean = true, + codegenDispatchConfigKey: Option[String] = None) private type CategoryNotes = Seq[ExprNotes] @@ -95,7 +102,15 @@ object GenerateDocs { case n: NativeOptInAvailable => n.hasConditionalNativeDefault case _ => false }, - codegenDispatchFallback = serde.isInstanceOf[CodegenDispatchFallback]) + codegenDispatchFallback = serde.isInstanceOf[CodegenDispatchFallback], + codegenDispatchEnabledByDefault = serde match { + case h: CodegenDispatchFallback => h.codegenDispatchEnabledByDefault + case _ => true + }, + codegenDispatchConfigKey = serde match { + case h: CodegenDispatchFallback => h.codegenDispatchConfigKey + case _ => None + }) } /** Build the documentation notes for a single aggregate expression serde. */ @@ -429,6 +444,13 @@ object GenerateDocs { " (run inside the Comet pipeline) by default." + s" Set `${n.nativeOptInConfigKey}=true` to explicitly select Comet's native" + " implementation, which has the following differences from Spark:\n\n" + } else if (n.codegenDispatchFallback && !n.codegenDispatchEnabledByDefault) { + val dispatchKey = + n.codegenDispatchConfigKey.getOrElse(s"spark.comet.expression.$name.codegen.enabled") + s"\nThe following incompatibilities cause `$name` to fall back to Spark by default." + + s" Set `$dispatchKey=true` to run Spark's code-generated implementation inside the" + + s" Comet pipeline. Set `${n.nativeOptInConfigKey}=true` to enable Comet's native" + + " implementation despite these differences.\n\n" } else if (n.nativeOptIn) { s"\nBy default, `$name` is evaluated in the JVM using Spark's own code-generated" + " implementation (run inside the Comet pipeline), which matches Spark exactly." + @@ -445,13 +467,20 @@ object GenerateDocs { } } if (n.unsupportedReasons.nonEmpty) { - val header = if (n.codegenDispatchFallback) { - "\nThe following cases have no native implementation and always run in the JVM using" + - " Spark's code-generated implementation (inside the Comet pipeline):\n\n" - } else { - "\nThe following cases are not supported by Comet and always fall back to Spark," + - " regardless of any `allowIncompatible` setting:\n\n" - } + val header = + if (n.codegenDispatchFallback && n.codegenDispatchEnabledByDefault) { + "\nThe following cases have no native implementation and always run in the JVM using" + + " Spark's code-generated implementation (inside the Comet pipeline):\n\n" + } else if (n.codegenDispatchFallback) { + val dispatchKey = n.codegenDispatchConfigKey.getOrElse( + s"spark.comet.expression.$name.codegen.enabled") + "\nThe following cases have no native implementation and fall back to Spark by" + + s" default. Set `$dispatchKey=true` to run Spark's code-generated implementation" + + " inside the Comet pipeline:\n\n" + } else { + "\nThe following cases are not supported by Comet and always fall back to Spark," + + " regardless of any `allowIncompatible` setting:\n\n" + } output.append(header) for (reason <- n.unsupportedReasons) { output.append(s"- $reason\n") diff --git a/spark/src/main/scala/org/apache/comet/serde/CometExpressionSerde.scala b/spark/src/main/scala/org/apache/comet/serde/CometExpressionSerde.scala index 206bd7c8207..98b1f4ab022 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometExpressionSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometExpressionSerde.scala @@ -114,6 +114,10 @@ trait CometExpressionSerde[T <: Expression] { * for the case where the dispatcher itself cannot handle the expression (e.g. the global codegen * flag is off, or the kernel rejects the bound tree). * + * A serde may override [[isCodegenDispatchEnabled]] so that this mixin is a no-op until the user + * opts in. When that method returns `false`, `Unsupported` and (non-opt-in) `Incompatible` fall + * back to Spark the same way a serde without this mixin would. + * * Contract for `Unsupported` reasons on a `CodegenDispatchFallback` serde: the case must be * something `Expression.doGenCode` can compile. If you mark something `Unsupported` because Spark * also rejects it, that is fine -- the dispatcher will surface the same error Spark would have. @@ -122,7 +126,27 @@ trait CometExpressionSerde[T <: Expression] { * dispatcher. Every other `Incompatible` expression falls back to Spark, and every other * `Unsupported` expression falls back to Spark. */ -trait CodegenDispatchFallback extends NativeOptInAvailable { self: CometExpressionSerde[_] => } +trait CodegenDispatchFallback extends NativeOptInAvailable { self: CometExpressionSerde[_] => + + /** + * Whether `Unsupported` / non-opt-in `Incompatible` cases should run through the JVM codegen + * dispatcher. Defaults to true. Override to gate the dispatcher behind a config; when false, + * those cases fall back to Spark. + */ + def isCodegenDispatchEnabled: Boolean = true + + /** + * Whether [[isCodegenDispatchEnabled]] is true at the config default. Documentation metadata + * only; does not affect runtime routing. + */ + def codegenDispatchEnabledByDefault: Boolean = true + + /** + * Config key the user sets to enable codegen dispatch when [[codegenDispatchEnabledByDefault]] + * is false. Documentation metadata only. + */ + def codegenDispatchConfigKey: Option[String] = None +} /** * Marker for serdes that have a native implementation the user can opt into. Normally these diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index eced2d7fdf9..7f7c16f884f 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -1245,18 +1245,19 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { } /** - * If `handler` is a `CodegenDispatchFallback`, run `expr` through the JVM codegen dispatcher - * and return `Some((handler, proto))` on success; otherwise return `None`. Shared by the - * `Unsupported` and (non-opt-in) `Incompatible` arms of `exprToProtoInternal` so they don't - * each inline the same pattern match. Returning the matched handler lets the `Incompatible` arm - * reach `nativeOptInConfigKeyOverride` without re-pattern-matching the same value. + * If `handler` is a `CodegenDispatchFallback` whose dispatcher is enabled, run `expr` through + * the JVM codegen dispatcher and return `Some((handler, proto))` on success; otherwise return + * `None`. Shared by the `Unsupported` and (non-opt-in) `Incompatible` arms of + * `exprToProtoInternal` so they don't each inline the same pattern match. Returning the matched + * handler lets the `Incompatible` arm reach `nativeOptInConfigKeyOverride` without + * re-pattern-matching the same value. */ private def dispatchIfFallback( handler: CometExpressionSerde[_], expr: Expression, inputs: Seq[Attribute], binding: Boolean): Option[(CodegenDispatchFallback, Expr)] = handler match { - case h: CodegenDispatchFallback => + case h: CodegenDispatchFallback if h.isCodegenDispatchEnabled => CometScalaUDF.emitJvmCodegenDispatch(expr, inputs, binding).map(h -> _) case _ => None } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index fc2e38c8f2a..edf882eb92f 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -452,10 +452,12 @@ object CometShuffleExchangeExec case MapType(keyType, valueType, _) if nestedHashPartitioningEnabled => // Map entry order is not semantically meaningful, so two equal maps must hash alike. // Spark 4.0+ normalizes a map shuffle key by wrapping it in `mapsort(...)`, which is - // gated separately by CometMapSort. Scalar keys use native map_sort; other orderable key - // types can use Spark's generated MapSort code through the JVM dispatcher. Earlier Spark - // versions insert no such normalization, so Comet would hash physical entry order and - // could route equal maps differently. + // gated separately by CometMapSort. Scalar keys use native map_sort. Other orderable key + // types fall the shuffle back to Spark unless + // `spark.comet.expression.MapSort.codegen.enabled` is set, in which case Spark's generated + // MapSort code runs through the JVM dispatcher. Earlier Spark versions insert no such + // normalization, so Comet would hash physical entry order and could route equal maps + // differently. isSpark40Plus && supportedHashPartitioningDataType(keyType) && supportedHashPartitioningDataType(valueType) diff --git a/spark/src/main/spark-4.x/org/apache/comet/serde/CometMapSort.scala b/spark/src/main/spark-4.x/org/apache/comet/serde/CometMapSort.scala index 1a48d099196..b10c91b8af5 100644 --- a/spark/src/main/spark-4.x/org/apache/comet/serde/CometMapSort.scala +++ b/spark/src/main/spark-4.x/org/apache/comet/serde/CometMapSort.scala @@ -25,13 +25,24 @@ import org.apache.spark.sql.types.MapType import org.apache.comet.CometConf import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, scalarFunctionExprToProtoWithReturnType, supportedScalarSortElementType} -// Key types without a native implementation can still run Spark's generated code in-pipeline. -// Spark rejects collated-string map keys by default, but they reach MapSort when +// Key types without a native implementation can run Spark's generated code in-pipeline when +// spark.comet.expression.MapSort.codegen.enabled=true. That dispatcher is off by default: a +// matched microbenchmark of those shapes is slower than falling the enclosing operator back to +// Spark. Spark rejects collated-string map keys by default, but they reach MapSort when // spark.sql.collation.allowInMapKeys=true and the map is built from dispatcher-supported inputs; -// those expressions take the same Unsupported -> dispatcher route. A scan carrying a collated -// map schema may still be rejected independently by the scan's schema support checks. +// with the codegen setting on, those expressions take the same Unsupported -> dispatcher route. +// A scan carrying a collated map schema may still be rejected independently by the scan's schema +// support checks. object CometMapSort extends CometExpressionSerde[MapSort] with CodegenDispatchFallback { + override def isCodegenDispatchEnabled: Boolean = + CometConf.COMET_EXPRESSION_MAPSORT_CODEGEN_ENABLED.get() + + override def codegenDispatchEnabledByDefault: Boolean = false + + override def codegenDispatchConfigKey: Option[String] = + Some(CometConf.COMET_EXPRESSION_MAPSORT_CODEGEN_ENABLED.key) + override def getIncompatibleReasons(): Seq[String] = Seq( "MapSort on floating-point keys is not 100% compatible with Spark when " + diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index 36b3db009f4..02a83c2f7fe 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -41,6 +41,9 @@ class CometMapExpressionSuite extends CometTestBase with CometCodegenAssertions s"expected MapSort in optimized plan:\n$plan") } + private def withMapSortCodegen(f: => Unit): Unit = + withSQLConf(CometConf.COMET_EXPRESSION_MAPSORT_CODEGEN_ENABLED.key -> "true")(f) + test("read map[int, int] from parquet") { withTempPath { dir => @@ -255,50 +258,70 @@ class CometMapExpressionSuite extends CometTestBase with CometCodegenAssertions } } - test("mapsort routes array keys through codegen dispatcher") { + test("mapsort falls back by default for array keys without codegen opt-in") { assume(isSpark40Plus, "Spark 4.0 inserts MapSort for group-by on map keys") - withTable("t_map_sort_array_key") { - sql("CREATE TABLE t_map_sort_array_key (m MAP, INT>) USING parquet") - sql("""INSERT INTO t_map_sort_array_key VALUES + withTable("t_map_sort_array_key_default") { + sql("CREATE TABLE t_map_sort_array_key_default (m MAP, INT>) USING parquet") + sql("""INSERT INTO t_map_sort_array_key_default VALUES |(map(array(2, 1), 20, array(1, 2), 10)), - |(map(array(1, 2), 10, array(2, 1), 20)), - |(map(array(3), 30)), - |(NULL)""".stripMargin) - val df = sql("SELECT m, count(*) FROM t_map_sort_array_key GROUP BY m") + |(map(array(1, 2), 10, array(2, 1), 20))""".stripMargin) + val df = sql("SELECT m, count(*) FROM t_map_sort_array_key_default GROUP BY m") assertMapSortInPlan(df) - val (_, cometPlan) = assertCodegenRan { - checkSparkAnswer(df) + assertCodegenDidNotRun { + checkSparkAnswerAndFallbackReason(df, "has no native implementation") + } + } + } + + test("mapsort routes array keys through codegen dispatcher") { + assume(isSpark40Plus, "Spark 4.0 inserts MapSort for group-by on map keys") + withMapSortCodegen { + withTable("t_map_sort_array_key") { + sql("CREATE TABLE t_map_sort_array_key (m MAP, INT>) USING parquet") + sql("""INSERT INTO t_map_sort_array_key VALUES + |(map(array(2, 1), 20, array(1, 2), 10)), + |(map(array(1, 2), 10, array(2, 1), 20)), + |(map(array(3), 30)), + |(NULL)""".stripMargin) + val df = sql("SELECT m, count(*) FROM t_map_sort_array_key GROUP BY m") + + assertMapSortInPlan(df) + val (_, cometPlan) = assertCodegenRan { + checkSparkAnswer(df) + } + val dispatched = new ExtendedExplainInfo().getCodegenDispatchExpressions(cometPlan) + assert( + dispatched.contains("mapsort"), + s"expected mapsort on codegen dispatch path, got $dispatched in:\n$cometPlan") } - val dispatched = new ExtendedExplainInfo().getCodegenDispatchExpressions(cometPlan) - assert( - dispatched.contains("mapsort"), - s"expected mapsort on codegen dispatch path, got $dispatched in:\n$cometPlan") } } test("mapsort routes struct keys through codegen dispatcher") { assume(isSpark40Plus, "Spark 4.0 inserts MapSort for group-by on map keys") - withTable("t_map_sort_struct_key") { - sql("""CREATE TABLE t_map_sort_struct_key ( - | m MAP, INT>) USING parquet""".stripMargin) - sql("""INSERT INTO t_map_sort_struct_key VALUES - |(map(named_struct('a', 2, 'b', 'b'), 20, - | named_struct('a', 1, 'b', 'a'), 10)), - |(map(named_struct('a', 1, 'b', 'a'), 10, - | named_struct('a', 2, 'b', 'b'), 20)), - |(map(named_struct('a', 3, 'b', 'c'), 30)), - |(NULL)""".stripMargin) - val df = sql("SELECT m, count(*) FROM t_map_sort_struct_key GROUP BY m") + withMapSortCodegen { + withTable("t_map_sort_struct_key") { + sql("""CREATE TABLE t_map_sort_struct_key ( + | m MAP, INT>) USING parquet""".stripMargin) + sql("""INSERT INTO t_map_sort_struct_key VALUES + |(map(named_struct('a', 2, 'b', 'b'), 20, + | named_struct('a', 1, 'b', 'a'), 10)), + |(map(named_struct('a', 1, 'b', 'a'), 10, + | named_struct('a', 2, 'b', 'b'), 20)), + |(map(named_struct('a', 3, 'b', 'c'), 30)), + |(NULL)""".stripMargin) + val df = sql("SELECT m, count(*) FROM t_map_sort_struct_key GROUP BY m") - assertMapSortInPlan(df) - val (_, cometPlan) = assertCodegenRan { - checkSparkAnswer(df) + assertMapSortInPlan(df) + val (_, cometPlan) = assertCodegenRan { + checkSparkAnswer(df) + } + val dispatched = new ExtendedExplainInfo().getCodegenDispatchExpressions(cometPlan) + assert( + dispatched.contains("mapsort"), + s"expected mapsort on codegen dispatch path, got $dispatched in:\n$cometPlan") } - val dispatched = new ExtendedExplainInfo().getCodegenDispatchExpressions(cometPlan) - assert( - dispatched.contains("mapsort"), - s"expected mapsort on codegen dispatch path, got $dispatched in:\n$cometPlan") } } @@ -335,6 +358,7 @@ class CometMapExpressionSuite extends CometTestBase with CometCodegenAssertions |(map(array(1, 2), 10, array(2, 1), 20))""".stripMargin) withSQLConf( CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_EXPRESSION_MAPSORT_CODEGEN_ENABLED.key -> "true", CometConf.getExprEnabledConfigKey("MapSort") -> "false") { val df = sql("SELECT m, count(*) FROM t_map_sort_dispatch_disabled GROUP BY m") @@ -353,7 +377,8 @@ class CometMapExpressionSuite extends CometTestBase with CometCodegenAssertions assume(isSpark40Plus, "collated map keys and MapSort require Spark 4.0+") withSQLConf( "spark.sql.collation.allowInMapKeys" -> "true", - CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_EXPRESSION_MAPSORT_CODEGEN_ENABLED.key -> "true") { withTable("t_map_sort_collated_key") { sql( "CREATE TABLE t_map_sort_collated_key " + @@ -441,6 +466,7 @@ class CometMapExpressionSuite extends CometTestBase with CometCodegenAssertions withSQLConf( CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true", CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_EXPRESSION_MAPSORT_CODEGEN_ENABLED.key -> "true", "spark.sql.legacy.disableMapKeyNormalization" -> "true") { withTable("t_map_sort_fp_key") { sql("CREATE TABLE t_map_sort_fp_key (id INT, m MAP) USING parquet") diff --git a/spark/src/test/scala/org/apache/comet/GenerateDocsSuite.scala b/spark/src/test/scala/org/apache/comet/GenerateDocsSuite.scala index 98827af378b..539f7883f80 100644 --- a/spark/src/test/scala/org/apache/comet/GenerateDocsSuite.scala +++ b/spark/src/test/scala/org/apache/comet/GenerateDocsSuite.scala @@ -109,4 +109,32 @@ class GenerateDocsSuite extends AnyFunSuite { markdown.contains( "The following cases are not supported by Comet and always fall back to Spark,")) } + + test("codegen-dispatch fallback that is off by default describes Spark fallback") { + val dispatchKey = "spark.comet.expression.MapSort.codegen.enabled" + val markdown = GenerateDocs.renderExpressionCompatNotes( + Seq(GenerateDocs.ExprNotes( + "MapSort", + Seq.empty, + Seq("Floating-point incompatibility."), + Seq("No native implementation for nested keys."), + nativeOptIn = true, + nativeOptInConfigKey = "spark.comet.expression.MapSort.allowIncompatible", + conditionalNativeDefault = false, + codegenDispatchFallback = true, + codegenDispatchEnabledByDefault = false, + codegenDispatchConfigKey = Some(dispatchKey)))) + + assert( + markdown.contains( + "The following incompatibilities cause `MapSort` to fall back to Spark by default.")) + assert( + markdown.contains(s"Set `$dispatchKey=true` to run Spark's code-generated implementation")) + assert( + !markdown.contains( + "The following cases have no native implementation and always run in the JVM")) + assert( + markdown.contains( + "The following cases have no native implementation and fall back to Spark by default.")) + } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala index a25316f1726..1edcac571c6 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala @@ -38,6 +38,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper { protected val adaptiveExecutionEnabled: Boolean @@ -142,7 +143,9 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar test("columnar shuffle on array/struct map key/value") { // Spark 4.0+ normalizes maps used as shuffle keys with mapsort(...). Native map_sort only - // supports scalar keys, so array and struct keys use Spark codegen in the Comet pipeline. + // supports scalar keys, so array and struct keys fall the shuffle back to Spark by default. + // Complex values with a scalar key stay on native map_sort and remain a Comet shuffle. + val complexKeyShuffles = if (isSpark40Plus) 0 else 1 Seq("false", "true").foreach { execEnabled => Seq(10, 201).foreach { numPartitions => Seq("1.0", "10.0").foreach { ratio => @@ -155,7 +158,7 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar .repartition(numPartitions, $"_1", $"_2") .sortWithinPartitions($"_2") - checkShuffleAnswer(df, 1) + checkShuffleAnswer(df, complexKeyShuffles) } withParquetTable((0 until 50).map(i => (Map(i -> Seq(i, i + 1)), i + 1)), "tbl") { @@ -173,7 +176,7 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar .repartition(numPartitions, $"_1", $"_2") .sortWithinPartitions($"_2") - checkShuffleAnswer(df, 1) + checkShuffleAnswer(df, complexKeyShuffles) } withParquetTable((0 until 50).map(i => (Map(i -> ((i, i.toString))), i + 1)), "tbl") { @@ -190,6 +193,29 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar } } + test("columnar shuffle on array/struct map key uses MapSort codegen dispatcher") { + assume(isSpark40Plus, "map shuffle keys are only normalized with mapsort on Spark 4.0+") + withSQLConf(CometConf.COMET_EXPRESSION_MAPSORT_CODEGEN_ENABLED.key -> "true") { + withParquetTable((0 until 50).map(i => (Map(Seq(i, i + 1) -> i), i + 1)), "tbl") { + val df = sql("SELECT * FROM tbl") + .filter($"_2" > 10) + .repartition(10, $"_1", $"_2") + .sortWithinPartitions($"_2") + + checkShuffleAnswer(df, 1) + } + + withParquetTable((0 until 50).map(i => (Map((i, i.toString) -> i), i + 1)), "tbl") { + val df = sql("SELECT * FROM tbl") + .filter($"_2" > 10) + .repartition(10, $"_1", $"_2") + .sortWithinPartitions($"_2") + + checkShuffleAnswer(df, 1) + } + } + } + test("columnar shuffle on map array element") { Seq("false", "true").foreach { execEnabled => Seq(10, 201).foreach { numPartitions => diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index 974ac39fc32..67aec64d887 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -871,8 +871,8 @@ class CometNativeShuffleSuite // Map entry order carries no meaning, so equal maps must hash alike. Spark 4.0+ normalizes a // map shuffle key with `mapsort(...)`; earlier versions do not, so Comet must not hash a raw // map there. The gate therefore only admits map keys on Spark 4.0+, and only when the - // `mapsort` itself can stay in the Comet pipeline, either natively or through the JVM - // codegen dispatcher. + // `mapsort` itself can stay in the Comet pipeline: native scalar keys, or other orderable + // keys when `spark.comet.expression.MapSort.codegen.enabled` is on. withParquetTable((0 until 50).map(i => (i, Map(i % 7 -> (i % 5)))), "tbl") { val df = sql("SELECT * FROM tbl").repartition(10, $"_2").sortWithinPartitions($"_1") @@ -881,17 +881,16 @@ class CometNativeShuffleSuite } } - test("native shuffle on map hash partitioning key with non-scalar map key uses dispatcher") { - // Arrow's sort_to_indices only handles scalar keys, so Spark's MapSort.doGenCode performs the - // normalization through the JVM codegen dispatcher while the shuffle remains native. + test("native shuffle on map hash partitioning key with non-scalar map key falls back") { + // Arrow's sort_to_indices only handles scalar keys. Without MapSort codegen opt-in, Spark + // 4.0+ mapsort for a nested map key cannot stay in the Comet pipeline, so the shuffle falls + // back rather than hashing an unnormalized map. assume(isSpark40Plus, "map shuffle keys are only normalized with mapsort on Spark 4.0+") withNestedHashPartitioning { withParquetTable((0 until 50).map(i => (i, Map(Seq(i % 7) -> (i % 5)))), "tbl") { val df = sql("SELECT * FROM tbl").repartition(10, $"_2").sortWithinPartitions($"_1") - assertCodegenRan { - checkShuffleAnswer(df, 1) - } + checkShuffleAnswer(df, 0) } withParquetTable( @@ -899,8 +898,33 @@ class CometNativeShuffleSuite "tbl") { val df = sql("SELECT * FROM tbl").repartition(10, $"_2").sortWithinPartitions($"_1") - assertCodegenRan { - checkShuffleAnswer(df, 1) + checkShuffleAnswer(df, 0) + } + } + } + + test("native shuffle on map hash partitioning key with non-scalar map key uses dispatcher") { + // Arrow's sort_to_indices only handles scalar keys, so Spark's MapSort.doGenCode performs the + // normalization through the JVM codegen dispatcher while the shuffle remains native. + assume(isSpark40Plus, "map shuffle keys are only normalized with mapsort on Spark 4.0+") + withSQLConf(CometConf.COMET_EXPRESSION_MAPSORT_CODEGEN_ENABLED.key -> "true") { + withNestedHashPartitioning { + withParquetTable((0 until 50).map(i => (i, Map(Seq(i % 7) -> (i % 5)))), "tbl") { + val df = sql("SELECT * FROM tbl").repartition(10, $"_2").sortWithinPartitions($"_1") + + assertCodegenRan { + checkShuffleAnswer(df, 1) + } + } + + withParquetTable( + (0 until 50).map(i => (i, Map((i % 7, (i % 5).toString) -> (i % 3)))), + "tbl") { + val df = sql("SELECT * FROM tbl").repartition(10, $"_2").sortWithinPartitions($"_1") + + assertCodegenRan { + checkShuffleAnswer(df, 1) + } } } } @@ -965,9 +989,10 @@ class CometNativeShuffleSuite } if (isSpark40Plus) { - withTable("complex_map_partition_keys") { - withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - sql("""CREATE TABLE complex_map_partition_keys USING parquet AS + withSQLConf(CometConf.COMET_EXPRESSION_MAPSORT_CODEGEN_ENABLED.key -> "true") { + withTable("complex_map_partition_keys") { + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sql("""CREATE TABLE complex_map_partition_keys USING parquet AS |SELECT CAST(id AS INT) AS id, | CASE WHEN pmod(id, 2) = 0 THEN | map(array(CAST(pmod(id, 13) AS INT), 1), CAST(pmod(id, 17) AS INT), @@ -988,53 +1013,54 @@ class CometNativeShuffleSuite | CAST(pmod(id, 17) AS INT)) | END AS struct_map |FROM range(200)""".stripMargin) - } - - assert( - CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_NESTED_ENABLED - .get(spark.sessionState.conf), - "complex map-key parity requires nested hash partitioning to be enabled") - - Seq("array_map", "struct_map").foreach { key => - val query = - "SELECT id, spark_partition_id() AS pid FROM (" + - s"SELECT /*+ REPARTITION(10, $key) */ id, $key " + - "FROM complex_map_partition_keys)" - val cometDf = sql(query) - val cometRows = assertCodegenRan { - cometDf.collect().map(r => (r.getInt(0), r.getInt(1))).sorted } - val cometPlan = stripAQEPlan(cometDf.queryExecution.executedPlan) - val cometExchanges = collect(cometPlan) { case e: CometShuffleExchangeExec => e } - assert( - cometExchanges.size == 1, - s"$key expected one CometShuffleExchangeExec:\n$cometPlan") + assert( - new ExtendedExplainInfo() - .getCodegenDispatchExpressions(cometPlan) - .contains("mapsort"), - s"$key expected MapSort on the JVM dispatcher path:\n$cometPlan") + CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_NESTED_ENABLED + .get(spark.sessionState.conf), + "complex map-key parity requires nested hash partitioning to be enabled") + + Seq("array_map", "struct_map").foreach { key => + val query = + "SELECT id, spark_partition_id() AS pid FROM (" + + s"SELECT /*+ REPARTITION(10, $key) */ id, $key " + + "FROM complex_map_partition_keys)" + val cometDf = sql(query) + val cometRows = assertCodegenRan { + cometDf.collect().map(r => (r.getInt(0), r.getInt(1))).sorted + } + val cometPlan = stripAQEPlan(cometDf.queryExecution.executedPlan) + val cometExchanges = collect(cometPlan) { case e: CometShuffleExchangeExec => e } + assert( + cometExchanges.size == 1, + s"$key expected one CometShuffleExchangeExec:\n$cometPlan") + assert( + new ExtendedExplainInfo() + .getCodegenDispatchExpressions(cometPlan) + .contains("mapsort"), + s"$key expected MapSort on the JVM dispatcher path:\n$cometPlan") + + var sparkRows: Array[(Int, Int)] = Array.empty + var sparkPartitions = -1 + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + val sparkDf = sql(query) + sparkRows = sparkDf.collect().map(r => (r.getInt(0), r.getInt(1))).sorted + val sparkPlan = stripAQEPlan(sparkDf.queryExecution.executedPlan) + val sparkExchanges = collect(sparkPlan) { case e: ShuffleExchangeExec => e } + assert( + sparkExchanges.size == 1, + s"$key expected one Spark ShuffleExchangeExec:\n$sparkPlan") + sparkPartitions = sparkExchanges.head.outputPartitioning.numPartitions + } - var sparkRows: Array[(Int, Int)] = Array.empty - var sparkPartitions = -1 - withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val sparkDf = sql(query) - sparkRows = sparkDf.collect().map(r => (r.getInt(0), r.getInt(1))).sorted - val sparkPlan = stripAQEPlan(sparkDf.queryExecution.executedPlan) - val sparkExchanges = collect(sparkPlan) { case e: ShuffleExchangeExec => e } + assert(sparkRows.nonEmpty, s"Spark produced no rows for $key") + assert( + cometExchanges.head.outputPartitioning.numPartitions == sparkPartitions, + s"$key output partition count differs from Spark") assert( - sparkExchanges.size == 1, - s"$key expected one Spark ShuffleExchangeExec:\n$sparkPlan") - sparkPartitions = sparkExchanges.head.outputPartitioning.numPartitions + cometRows === sparkRows, + s"partition assignment differs from Spark for map key $key") } - - assert(sparkRows.nonEmpty, s"Spark produced no rows for $key") - assert( - cometExchanges.head.outputPartitioning.numPartitions == sparkPartitions, - s"$key output partition count differs from Spark") - assert( - cometRows === sparkRows, - s"partition assignment differs from Spark for map key $key") } } } @@ -1127,38 +1153,49 @@ class CometNativeShuffleSuite true)) cases.foreach { case (name, values, expectDispatch) => - val query = s"""SELECT id, spark_partition_id() AS pid FROM ( - | SELECT /*+ REPARTITION(10, m) */ * FROM VALUES - | $values AS t(id, m))""".stripMargin - val df = sql(query) - val rows = { - def collectRows(): Map[Int, Int] = - df.collect().map(r => (r.getInt(0), r.getInt(1))).toMap - if (expectDispatch) assertCodegenRan(collectRows()) else collectRows() - } - assert( - rows(1) == rows(2), - s"equal $name-keyed maps in different entry order must share a partition") - assert( - rows(3) == rows(4), - s"equal $name-keyed maps in different entry order must share a partition") + def runCase(): Unit = { + val query = s"""SELECT id, spark_partition_id() AS pid FROM ( + | SELECT /*+ REPARTITION(10, m) */ * FROM VALUES + | $values AS t(id, m))""".stripMargin + val df = sql(query) + val rows = { + def collectRows(): Map[Int, Int] = + df.collect().map(r => (r.getInt(0), r.getInt(1))).toMap + if (expectDispatch) assertCodegenRan(collectRows()) else collectRows() + } + assert( + rows(1) == rows(2), + s"equal $name-keyed maps in different entry order must share a partition") + assert( + rows(3) == rows(4), + s"equal $name-keyed maps in different entry order must share a partition") - val cometPlan = stripAQEPlan(df.queryExecution.executedPlan) - assert( - collect(cometPlan) { case e: CometShuffleExchangeExec => e }.nonEmpty, - s"$name-keyed map did not retain native shuffle:\n$cometPlan") - val dispatched = - new ExtendedExplainInfo().getCodegenDispatchExpressions(cometPlan).contains("mapsort") - assert( - dispatched == expectDispatch, - s"unexpected mapsort dispatch route for $name-keyed map:\n$cometPlan") + val cometPlan = stripAQEPlan(df.queryExecution.executedPlan) + assert( + collect(cometPlan) { case e: CometShuffleExchangeExec => e }.nonEmpty, + s"$name-keyed map did not retain native shuffle:\n$cometPlan") + val dispatched = + new ExtendedExplainInfo() + .getCodegenDispatchExpressions(cometPlan) + .contains("mapsort") + assert( + dispatched == expectDispatch, + s"unexpected mapsort dispatch route for $name-keyed map:\n$cometPlan") - var sparkRows: Map[Int, Int] = Map.empty - withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - sparkRows = sql(query).collect().map(r => (r.getInt(0), r.getInt(1))).toMap + var sparkRows: Map[Int, Int] = Map.empty + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkRows = sql(query).collect().map(r => (r.getInt(0), r.getInt(1))).toMap + } + assert(sparkRows.nonEmpty) + assert(rows == sparkRows, s"$name-keyed map partition assignment differs from Spark") + } + if (expectDispatch) { + withSQLConf(CometConf.COMET_EXPRESSION_MAPSORT_CODEGEN_ENABLED.key -> "true") { + runCase() + } + } else { + runCase() } - assert(sparkRows.nonEmpty) - assert(rows == sparkRows, s"$name-keyed map partition assignment differs from Spark") } } } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala index e13704efc7f..83dc7225899 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala @@ -38,15 +38,15 @@ import org.apache.comet.udf.codegen.CometScalaUDFCodegen * Matched benchmark for the two routes Spark 4.x can take for `MapSort` shapes that Comet cannot * sort natively: * - * - with `spark.comet.expression.MapSort.enabled=false`, the enclosing projection or shuffle - * falls back to Spark; and - * - with MapSort enabled, Spark's `MapSort.doGenCode` executes inside the Comet pipeline. + * - with `spark.comet.expression.MapSort.codegen.enabled=false` (the default), the enclosing + * projection or shuffle falls back to Spark; and + * - with that setting true, Spark's `MapSort.doGenCode` executes inside the Comet pipeline. * - * Every pair reads the same Parquet data and differs only in - * `spark.comet.expression.MapSort.enabled`; the global codegen dispatcher remains enabled in both - * arms. Array and struct cases vary map size independently from nested-key width. Strict - * floating-point cases include NaN and both signed zeros in the same map. Input maps are written - * in reverse key order and one row in 64 has a NULL map. + * Every pair keeps `spark.comet.expression.MapSort.enabled=true` and the global codegen + * dispatcher enabled, and differs only in `spark.comet.expression.MapSort.codegen.enabled`. Array + * and struct cases vary map size independently from nested-key width. Strict floating-point cases + * include NaN and both signed zeros in the same map. Input maps are written in reverse key order + * and one row in 64 has a NULL map. * * Spark 4.0 and 4.1 only insert `MapSort` for grouping and repartition expressions; * `try_element_at` itself does not insert one. To measure a projection without also timing an @@ -642,7 +642,8 @@ object CometMapSortBenchmark extends CometBenchmarkBase { CometConf.COMET_ENABLED.key -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true", CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", - CometConf.getExprEnabledConfigKey("MapSort") -> dispatch.toString, + CometConf.getExprEnabledConfigKey("MapSort") -> "true", + CometConf.COMET_EXPRESSION_MAPSORT_CODEGEN_ENABLED.key -> dispatch.toString, CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> (shape.family == StrictDoubleKey).toString, "spark.sql.legacy.disableMapKeyNormalization" -> @@ -784,7 +785,9 @@ object CometMapSortBenchmark extends CometBenchmarkBase { emit(s"Mode: $Mode; steady-state case order: $CaseOrder") emit("Shuffle mode: native; AQE: disabled; nested hash partitioning: enabled") emit(s"Global dispatcher enabled: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=true") - emit(s"Only matched-arm difference: ${CometConf.getExprEnabledConfigKey("MapSort")}") + emit(s"${CometConf.getExprEnabledConfigKey("MapSort")}=true in both arms") + emit( + s"Only matched-arm difference: ${CometConf.COMET_EXPRESSION_MAPSORT_CODEGEN_ENABLED.key}") emit(s"Pinned: ${CometConf.getExprAllowIncompatConfigKey("MapSort")}=false") emit("Projection note: Spark 4.0/4.1 do not insert MapSort for try_element_at. This suite") emit(" executes the optimizer-inserted grouping MapSort Project alone; no aggregate is run.") From 04ba768f15dd2309b0d183a2146d0170eb7cd2e6 Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Tue, 22 Sep 2026 11:13:56 +0800 Subject: [PATCH 7/8] test: capture the target Parquet insert plan, not a delayed source INSERT --- .../parquet/CometParquetWriterSuite.scala | 28 +++++--- .../parquet/CometParquetWriterTestBase.scala | 67 ++++++++++++++++--- 2 files changed, 76 insertions(+), 19 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala index c46524e2917..9adbbb970fe 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -34,7 +34,7 @@ import org.apache.parquet.hadoop.util.HadoopInputFile import org.apache.parquet.schema.{MessageType, Type} import org.apache.spark.internal.io.FileCommitProtocol import org.apache.spark.sql.{AnalysisException, DataFrame, Row, SaveMode} -import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.{InternalRow, TableIdentifier} import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeScanExec, CometScanExec, CometWriteFilesExec} import org.apache.spark.sql.execution.{FileSourceScanExec, SparkPlan} import org.apache.spark.sql.execution.datasources.{BasicWriteTaskStats, SQLHadoopMapReduceCommitProtocol, WriteTaskStats, WriteTaskStatsTracker} @@ -1183,10 +1183,16 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } withNativeWriter { // Assert the write itself went native: a fallback to Spark's writer would make the - // read-back pass for the wrong reason. - assertHasCometNativeWriteExec( - captureWritePlan( - sql("INSERT INTO comet_write_target SELECT id, name FROM comet_write_source"))) + // read-back pass for the wrong reason. Match comet_write_target's location so a late + // source INSERT VALUES event is not the plan under test. + val targetLocation = spark.sessionState.catalog + .getTableMetadata(TableIdentifier("comet_write_target")) + .storage + .locationUri + .getOrElse(fail("comet_write_target has no storage location")) + assertHasCometNativeWriteExec(captureWritePlan(new Path(targetLocation)) { + sql("INSERT INTO comet_write_target SELECT id, name FROM comet_write_source") + }) } checkAnswer(spark.table("comet_write_target"), Row(1L, "a") :: Row(2L, "b") :: Nil) } @@ -1308,9 +1314,15 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } withNativeWriter { - assertHasCometNativeWriteExec( - captureWritePlan( - sql("INSERT INTO comet_rename_target SELECT id + 1, s FROM comet_rename_source"))) + // Match the target location so a late source INSERT VALUES event is not captured. + val targetLocation = spark.sessionState.catalog + .getTableMetadata(TableIdentifier("comet_rename_target")) + .storage + .locationUri + .getOrElse(fail("comet_rename_target has no storage location")) + assertHasCometNativeWriteExec(captureWritePlan(new Path(targetLocation)) { + sql("INSERT INTO comet_rename_target SELECT id + 1, s FROM comet_rename_source") + }) } // Read the names out of the file itself: the catalog would report the target's schema diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala index f3e59e75101..4dfdf76687a 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala @@ -19,10 +19,14 @@ package org.apache.comet.parquet +import java.util.concurrent.atomic.AtomicReference + +import org.apache.hadoop.fs.Path import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.comet.{CometNativeWriteExec, CometWriteFilesExec} import org.apache.spark.sql.execution.{QueryExecution, SparkPlan} import org.apache.spark.sql.execution.command.DataWritingCommandExec +import org.apache.spark.sql.execution.datasources.InsertIntoHadoopFsRelationCommand import org.apache.spark.sql.internal.SQLConf import org.apache.comet.{CometConf, CometExplainInfo} @@ -57,16 +61,29 @@ abstract class CometParquetWriterTestBase extends CometTestBase { * The captured execution plan */ protected def captureWritePlan(writeOp: String => Unit, outputPath: String): SparkPlan = - captureWritePlan(writeOp(outputPath)) + captureWritePlan(new Path(outputPath))(writeOp(outputPath)) - /** As above, for a write that names its own target (an `INSERT INTO`, for example). */ - protected def captureWritePlan(writeOp: => Unit): SparkPlan = { - var capturedPlan: Option[QueryExecution] = None + /** + * Captures the execution plan of the write to `expectedOutputPath`. + * + * Only an [[InsertIntoHadoopFsRelationCommand]] whose output path is that target is accepted. + * The first `save` or command event is not enough: a preceding `INSERT VALUES` can arrive after + * this listener is registered, and that event belongs to a different write. + * + * @param expectedOutputPath + * output path of the write to capture + * @param writeOp + * the write to execute + * @return + * the captured execution plan + */ + protected def captureWritePlan(expectedOutputPath: Path)(writeOp: => Unit): SparkPlan = { + val capturedPlan = new AtomicReference[QueryExecution]() val listener = new org.apache.spark.sql.util.QueryExecutionListener { override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = { - if (funcName == "save" || funcName.contains("command")) { - capturedPlan = Some(qe) + if (writesTo(qe, expectedOutputPath)) { + capturedPlan.set(qe) } } @@ -81,27 +98,55 @@ abstract class CometParquetWriterTestBase extends CometTestBase { try { writeOp - // Wait for listener to be called with timeout + // The matching callback can still be queued after the write returns. val maxWaitTimeMs = 15000 val checkIntervalMs = 100 val maxIterations = maxWaitTimeMs / checkIntervalMs var iterations = 0 - while (capturedPlan.isEmpty && iterations < maxIterations) { + while (capturedPlan.get() == null && iterations < maxIterations) { Thread.sleep(checkIntervalMs) iterations += 1 } + val captured = Option(capturedPlan.get()) assert( - capturedPlan.isDefined, - s"Listener was not called within ${maxWaitTimeMs}ms - no execution plan captured") + captured.isDefined, + s"No write to $expectedOutputPath was captured within ${maxWaitTimeMs}ms") - stripAQEPlan(capturedPlan.get.executedPlan) + stripAQEPlan(captured.get.executedPlan) } finally { spark.listenerManager.unregister(listener) } } + /** Whether `qe` is an insert into `expectedOutputPath`. */ + private def writesTo(qe: QueryExecution, expectedOutputPath: Path): Boolean = + insertOutputPath(qe).exists(actual => sameOutputPath(actual, expectedOutputPath)) + + /** + * Output path of the insert command, from the optimized logical plan. On Spark 3.x the physical + * plan may already be a [[CometNativeWriteExec]], so the command is read from the logical plan + * first and from [[DataWritingCommandExec]] only as a fallback. + */ + private def insertOutputPath(qe: QueryExecution): Option[Path] = { + val fromOptimized = qe.optimizedPlan.collectFirst { + case cmd: InsertIntoHadoopFsRelationCommand => cmd.outputPath + } + fromOptimized.orElse { + qe.executedPlan.collectFirst { case exec: DataWritingCommandExec => + exec.cmd + } match { + case Some(cmd: InsertIntoHadoopFsRelationCommand) => Some(cmd.outputPath) + case _ => None + } + } + } + + /** Hadoop `Path` equality, then equality after re-parsing each path's string form. */ + private def sameOutputPath(actual: Path, expected: Path): Boolean = + actual == expected || new Path(actual.toString) == new Path(expected.toString) + /** * The operator that carries a native write, which differs by Spark version: on 4.0+ Comet * replaces only `WriteFilesExec` with [[CometWriteFilesExec]] and leaves Spark's write From 2c868d867421780b0dcc13ba171ee3a310efa125 Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Tue, 22 Sep 2026 15:25:06 +0800 Subject: [PATCH 8/8] test: qualify Parquet output paths before matching write plans Spark qualifies file-format write paths when building InsertIntoHadoopFsRelationCommand, while the test helper may receive an unqualified temporary path. Qualify both paths through Hadoop FileSystem before comparing them so captureWritePlan still filters delayed unrelated writes without missing the target write. Co-authored-by: Cursor --- .../parquet/CometParquetWriterTestBase.scala | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala index 4dfdf76687a..67c1f95833b 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala @@ -143,9 +143,18 @@ abstract class CometParquetWriterTestBase extends CometTestBase { } } - /** Hadoop `Path` equality, then equality after re-parsing each path's string form. */ - private def sameOutputPath(actual: Path, expected: Path): Boolean = - actual == expected || new Path(actual.toString) == new Path(expected.toString) + /** + * Spark qualifies file-format write paths when building [[InsertIntoHadoopFsRelationCommand]], + * while the helper may receive an unqualified temporary path. Re-parsing `Path.toString` does + * not drop a `file:` scheme, so both sides are qualified through Hadoop `FileSystem` before + * comparison. + */ + private def sameOutputPath(actual: Path, expected: Path): Boolean = { + val conf = spark.sessionState.newHadoopConf() + val actualFs = actual.getFileSystem(conf) + val expectedFs = expected.getFileSystem(conf) + actualFs.makeQualified(actual) == expectedFs.makeQualified(expected) + } /** * The operator that carries a native write, which differs by Spark version: on 4.0+ Comet