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..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,14 +21,23 @@ 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, so map shuffle and group-by-on-map stay on Comet under Spark 4.0. - -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`. +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. 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..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,12 @@ 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 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] = Seq( @@ -33,12 +38,14 @@ object CometMapSort extends CometExpressionSerde[MapSort] { 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") 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..36b3db009f4 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 @@ -30,8 +30,16 @@ 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 { +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 +255,240 @@ 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 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 -> "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, + "Expression support is disabled. Set " + + s"${CometConf.getExprEnabledConfigKey("MapSort")}=true to enable it.") + } + } + } + } + + 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", + 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 (id INT, m MAP) USING parquet") + sql("""INSERT INTO t_map_sort_fp_key VALUES + |(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) + 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") + 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.") + } + } + } + } + } + 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..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,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, ExtendedExplainInfo, 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: _*) { @@ -434,7 +437,7 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper 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 => @@ -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) + } + } } } @@ -941,8 +958,6 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper "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 @@ -959,6 +974,81 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper } } + 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) { @@ -1017,23 +1107,69 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper // 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 new file mode 100644 index 00000000000..e13704efc7f --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala @@ -0,0 +1,812 @@ +/* + * 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, 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 + +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 `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.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 + * 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. + * + * 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 \ + * 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) + | 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 >= 3, "mapSize must leave room for NaN, -0.0, 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, dispatcher.plan) + assertRoute(shape, "projection", dispatch = false, fallback.plan) + 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, 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}", + 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(" | ") + + /** + * 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): PreparedQuery = { + var prepared: PreparedQuery = null + withSQLConf(configs(shape, dispatch, shuffle): _*) { + 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 + } + + /** + * Executes and fully consumes an already-prepared physical plan. Does not construct a + * `QueryExecution`. + */ + private def executePreparedQuery( + shape: Shape, + dispatch: Boolean, + shuffle: Boolean, + prepared: PreparedQuery, + plan: SparkPlan): Unit = + withSQLConf(configs(shape, dispatch, shuffle): _*) { + 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: 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" => + 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 -> "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", + 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 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( + executePreparedQuery( + shape, + dispatch = dispatch, + shuffle = shuffle, + prepared = prepared, + plan = prepared.plan)) + val stats = CometScalaUDFCodegen.stats() + 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(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.") + } + } + + 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) + 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 = + 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"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.") + 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")) +}