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..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 @@ -21,14 +21,25 @@ 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 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 d371f1ba44c..c972be5a12f 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 43d65fd418e..b708ef784cd 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -1249,18 +1249,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 2a539643755..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,9 +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 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 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 4fad289a080..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,7 +25,23 @@ 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 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; +// 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( @@ -33,12 +49,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..02a83c2f7fe 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,19 @@ 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") + } + + private def withMapSortCodegen(f: => Unit): Unit = + withSQLConf(CometConf.COMET_EXPRESSION_MAPSORT_CODEGEN_ENABLED.key -> "true")(f) test("read map[int, int] from parquet") { @@ -247,6 +258,263 @@ class CometMapExpressionSuite extends CometTestBase { } } + 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_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))""".stripMargin) + val df = sql("SELECT m, count(*) FROM t_map_sort_array_key_default GROUP BY m") + + assertMapSortInPlan(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") + } + } + } + + test("mapsort routes struct keys through codegen dispatcher") { + assume(isSpark40Plus, "Spark 4.0 inserts MapSort for group-by on map keys") + 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) + } + 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.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") + + 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", + CometConf.COMET_EXPRESSION_MAPSORT_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", + 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") + 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/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 a44aac42b3a..1edcac571c6 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala @@ -142,9 +142,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(...). 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. + // Spark 4.0+ normalizes maps used as shuffle keys with mapsort(...). Native map_sort only + // 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 => @@ -193,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 353ee66d4d1..67aec64d887 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: _*) { @@ -424,7 +427,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 => @@ -868,7 +871,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: 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") @@ -878,14 +882,51 @@ 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. + // 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+") - 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) + checkShuffleAnswer(df, 0) + } + + 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") + + 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) + } + } + } } } @@ -931,8 +972,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 @@ -949,6 +988,83 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper } } + if (isSpark40Plus) { + 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), + | 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) { @@ -1007,23 +1123,80 @@ 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") - - 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 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) => + 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") + + 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, "map key partition assignment differs from Spark") } } } 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..67c1f95833b 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,64 @@ 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 + } + } + } + + /** + * 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 * replaces only `WriteFilesExec` with [[CometWriteFilesExec]] and leaves Spark's write 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..83dc7225899 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapSortBenchmark.scala @@ -0,0 +1,815 @@ +/* + * 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.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 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 + * 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") -> "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" -> + (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"${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.") + 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")) +}