From 49a7196910ab210095a9393dc46c0e325d8e5cc7 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sun, 6 Sep 2026 18:01:17 -0700 Subject: [PATCH 1/2] bench: measure nested hash partitioning keys in native shuffle Nothing covered the shapes admitted by `spark.comet.shuffle.native.partitioning.hash.nested.enabled`. `shuffleArrayBenchmark` and `shuffleStructBenchmark` do repartition on a nested column but have no native-shuffle case, because the gate rejected nested keys when they were written, and `shuffleDeeplyNestedBenchmark` calls `repartition(n)` with no key, which is round robin rather than hash partitioning. Adds `shuffleNestedHashKeyBenchmark` with a native-shuffle case that enables the config, so it measures the native hashing path rather than a silent fallback. The shapes separate the two code paths in the native Murmur3 kernel: a list whose elements are primitives is vectorized, while a list whose elements are nested falls through to `hash_list_array!`, which slices a one-element array and re-enters `create_murmur3_hashes` for every element. Measured on an Apple M4 Max, local[5], 1M rows (best time, relative to Spark doing the whole shuffle): struct 5 partitions 79ms -> 41ms 1.9x struct 201 partitions 126ms -> 89ms 1.4x array 5 partitions 76ms -> 52ms 1.5x array 201 partitions 145ms -> 101ms 1.4x struct, string> 5 partitions 100ms -> 68ms 1.5x array> 5 partitions 164ms -> 561ms 0.3x array> 201 partitions 250ms -> 647ms 0.4x struct, int> 5 partitions 121ms -> 97ms 1.2x struct, int> 201 partitions 160ms -> 155ms 1.0x The native case is not silently falling back to Spark's shuffle. Rerunning the `array>` case at 5 partitions separates the three Comet cases, at 761ms (Spark shuffle), 1059ms (JVM shuffle) and 2340ms (native shuffle), so the native path is distinct and is the slow one. Absolute times move with machine load, but the ordering and the set of shapes that win or lose reproduce. So `array>` is the one shape that loses, at 0.3-0.4x of leaving the shuffle to Spark, which is what the config defaults to off for. Every other shape measured is faster natively. That is the fallback path tracked in #5777, so these numbers are the before side of that comparison rather than a standing property of the kernel. A map nested in a struct is not affected, because the kernel specializes common scalar key/value pairs. Co-authored-by: Claude Code --- .../sql/benchmark/CometShuffleBenchmark.scala | 94 +++++++++++++++++++ 1 file changed, 94 insertions(+) diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala index 9f5aaef306a..32bc5907301 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala @@ -475,6 +475,81 @@ object CometShuffleBenchmark extends CometBenchmarkBase { } } + /** + * Nested hash partitioning keys, which native shuffle admits only when + * `spark.comet.shuffle.native.partitioning.hash.nested.enabled` is on. + * + * The shapes are chosen to separate the two code paths in the native Murmur3 kernel. + * `hash_list_with_primitive_elements!` vectorizes a list whose elements are primitives, but a + * list whose elements are themselves nested falls through to `hash_list_array!`, which slices a + * one-element array and re-enters `create_murmur3_hashes` per element; a struct additionally + * copies its column vector per call. `array>` is the shape that pays that, and it is + * the one where native hashing loses to letting Spark do the whole shuffle. A map nested in a + * struct does not: `hash_funcs/utils.rs` specializes common scalar key/value pairs, so it stays + * on a batched path. + * + * Each case is compared against Spark doing the whole shuffle, which is what the config's + * default trades against: if the native path is slower, falling back is the better default. + */ + def shuffleNestedHashKeyBenchmark( + name: String, + keyExpr: String, + values: Int, + partitionNum: Int): Unit = { + val benchmark = + microBenchmark(s"Nested hash key: $name ($partitionNum Partition)", values) + + withTempPath { dir => + withTempTable("parquetV1Table") { + // `tbl`'s `value` spans the full Long range, so a direct cast to INT overflows under ANSI + // mode. `pmod` keeps the key varied (a constant would hash every row alike, which would + // not measure partitioning at all) while staying in range. + prepareTable(dir, spark.sql(s"SELECT CAST(pmod(value, 1000000) AS INT) AS c1 FROM $tbl")) + val query = s"SELECT $keyExpr AS k, c1 FROM parquetV1Table" + + benchmark.addCase("Spark") { _ => + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.sql(query).repartition(partitionNum, Column("k")).noop() + } + } + + benchmark.addCase("Comet (Spark Shuffle)") { _ => + withSQLConf( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_ENABLED.key -> "false") { + spark.sql(query).repartition(partitionNum, Column("k")).noop() + } + } + + benchmark.addCase("Comet (JVM Shuffle)") { _ => + withSQLConf( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + spark.sql(query).repartition(partitionNum, Column("k")).noop() + } + } + + // Nested keys are rejected by the native gate unless the config is on, so without it this + // case would silently measure a fallback rather than the native hashing path. + benchmark.addCase("Comet (Native Shuffle)") { _ => + withSQLConf( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_NESTED_ENABLED.key -> "true") { + spark.sql(query).repartition(partitionNum, Column("k")).noop() + } + } + + benchmark.run() + } + } + } + override def runCometBenchmark(mainArgs: Array[String]): Unit = { // nested type shuffle @@ -492,6 +567,25 @@ object CometShuffleBenchmark extends CometBenchmarkBase { } } + runBenchmarkWithTable("Nested hash partitioning key", 1024 * 1024 * 1) { v => + // Shapes whose leaves are primitives take the vectorized element path; the last two force + // the per-element path in the native kernel. + val shapes = Seq( + "struct" -> "named_struct('a', c1, 'b', CAST(c1 AS STRING))", + "array" -> "ARRAY_REPEAT(c1, 10)", + "struct, string>" -> + "named_struct('a', ARRAY_REPEAT(c1, 10), 'b', CAST(c1 AS STRING))", + "array>" -> + "ARRAY_REPEAT(named_struct('a', c1, 'b', CAST(c1 AS STRING)), 10)", + "struct, int>" -> + "named_struct('m', MAP(CAST(c1 AS STRING), c1), 'i', c1)") + shapes.foreach { case (name, keyExpr) => + Seq(5, manyPartitions).foreach { partitionNum => + shuffleNestedHashKeyBenchmark(name, keyExpr, v, partitionNum) + } + } + } + runBenchmarkWithTable("Shuffle on array", 1024 * 1024 * 1) { v => benchmarkTypes.foreach { dataType => Seq(5, manyPartitions).foreach { partitionNum => From 8d9eede019d09e18c85c18f8b383d3e2c2cd1ede Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Wed, 9 Sep 2026 11:04:16 -0700 Subject: [PATCH 2/2] bench: verify nested shuffle admission and expand map coverage Skip native map-key cases before Spark 4.0 and verify the native exchange before timing. Compare against the JVM shuffle baseline, add variable-size maps in both input orders, and provide a nested-hash-only entry point. Validated Spark 4.1.3 with 18 groups / 72 benchmark cases and native plan checks. Spark 3.5.9 clean build and 14-case smoke run confirm native struct admission and explicit native map skips at 5 and 201 partitions. --- .../sql/benchmark/CometShuffleBenchmark.scala | 114 ++++++++++++------ 1 file changed, 80 insertions(+), 34 deletions(-) diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala index 32bc5907301..64d9d0237af 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala @@ -27,17 +27,20 @@ import scala.util.Random import org.apache.spark.SparkConf import org.apache.spark.benchmark.Benchmark import org.apache.spark.sql.{Column, SaveMode, SparkSession} +import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, SchemaGenOptions} // spotless:off /** * Benchmark to measure Comet shuffle performance. To run this benchmark: * `SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometShuffleBenchmark` + * Add `-- --nested-hash-only` to run just the nested hash key cases. * Results will be written to "spark/benchmarks/CometShuffleBenchmark-**results.txt". */ // spotless:on @@ -479,17 +482,13 @@ object CometShuffleBenchmark extends CometBenchmarkBase { * Nested hash partitioning keys, which native shuffle admits only when * `spark.comet.shuffle.native.partitioning.hash.nested.enabled` is on. * - * The shapes are chosen to separate the two code paths in the native Murmur3 kernel. - * `hash_list_with_primitive_elements!` vectorizes a list whose elements are primitives, but a - * list whose elements are themselves nested falls through to `hash_list_array!`, which slices a - * one-element array and re-enters `create_murmur3_hashes` per element; a struct additionally - * copies its column vector per call. `array>` is the shape that pays that, and it is - * the one where native hashing loses to letting Spark do the whole shuffle. A map nested in a - * struct does not: `hash_funcs/utils.rs` specializes common scalar key/value pairs, so it stays - * on a batched path. + * Primitive arrays use the typed element path; arrays of structs exercise recursive hashing. + * Map cases include a singleton control and variable cardinalities in both input key orders, + * covering normalization as well as the specialized scalar key/value hash loop. * - * Each case is compared against Spark doing the whole shuffle, which is what the config's - * default trades against: if the native path is slower, falling back is the better default. + * Compare native with Comet JVM shuffle to evaluate the default `auto` mode with nested hashing + * disabled. The all-Spark arm also changes scan and projection execution. These are end-to-end + * shuffle measurements, not isolated hash-kernel timings. */ def shuffleNestedHashKeyBenchmark( name: String, @@ -532,16 +531,44 @@ object CometShuffleBenchmark extends CometBenchmarkBase { } } - // Nested keys are rejected by the native gate unless the config is on, so without it this - // case would silently measure a fallback rather than the native hashing path. - benchmark.addCase("Comet (Native Shuffle)") { _ => - withSQLConf( + def containsMap(dataType: DataType): Boolean = dataType match { + case _: MapType => true + case ArrayType(elementType, _) => containsMap(elementType) + case StructType(fields) => fields.exists(f => containsMap(f.dataType)) + case _ => false + } + + // Spark 3.x does not normalize map partitioning keys for native hashing. + if (containsMap(spark.sql(query).schema("k").dataType) && !isSpark40Plus) { + val message = s"Skipping native shuffle for $name: map keys require Spark 4.0+" + benchmark.out.println(message) + } else { + val nativeConfigs = Seq( CometConf.COMET_ENABLED.key -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true", CometConf.COMET_SHUFFLE_ENABLED.key -> "true", CometConf.COMET_SHUFFLE_MODE.key -> "native", - CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_NESTED_ENABLED.key -> "true") { - spark.sql(query).repartition(partitionNum, Column("k")).noop() + CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_NESTED_ENABLED.key -> "true") + // Check outside the timer: enabling the gate alone does not prove native admission. + withSQLConf(nativeConfigs: _*) { + val plan = + spark.sql(query).repartition(partitionNum, Column("k")).queryExecution.executedPlan + val nativeExchanges = collect(plan) { + case exchange: CometShuffleExchangeExec + if exchange.shuffleType == CometNativeShuffle => + exchange + } + require( + nativeExchanges.size == 1, + s"Expected one native shuffle for $name, found ${nativeExchanges.size}:\n$plan") + benchmark.out.println( + s"Verified native exchange for $name ($partitionNum partitions):") + benchmark.out.println(plan.treeString) + } + benchmark.addCase("Comet (Native Shuffle)") { _ => + withSQLConf(nativeConfigs: _*) { + spark.sql(query).repartition(partitionNum, Column("k")).noop() + } } } @@ -550,7 +577,43 @@ object CometShuffleBenchmark extends CometBenchmarkBase { } } + private def runNestedHashKeyBenchmarks(): Unit = { + runBenchmarkWithTable("Nested hash partitioning key", 1024 * 1024 * 1) { v => + val shapes = Seq( + "struct" -> "named_struct('a', c1, 'b', CAST(c1 AS STRING))", + "array" -> "ARRAY_REPEAT(c1, 10)", + "struct, string>" -> + "named_struct('a', ARRAY_REPEAT(c1, 10), 'b', CAST(c1 AS STRING))", + "array>" -> + "ARRAY_REPEAT(named_struct('a', c1, 'b', CAST(c1 AS STRING)), 10)", + "struct, int>" -> + "named_struct('m', MAP(CAST(c1 AS STRING), c1), 'i', c1)") + // Distinct keys, variable entry counts, and opposite input orders exercise map sorting. + val mapShapes = for { + maxEntries <- Seq(10, 50) + reverse <- Seq(false, true) + } yield { + val indices = s"sequence(1, 2 + pmod(c1, ${maxEntries - 1}))" + val ordered = if (reverse) s"reverse($indices)" else indices + val map = s"map_from_arrays(transform($ordered, x -> CAST(c1 + x AS STRING)), " + + s"transform($ordered, x -> c1 + x))" + val order = if (reverse) "reversed" else "forward" + s"struct, int> (2-$maxEntries entries, $order)" -> + s"named_struct('m', $map, 'i', c1)" + } + (shapes ++ mapShapes).foreach { case (name, keyExpr) => + Seq(5, manyPartitions).foreach { partitionNum => + shuffleNestedHashKeyBenchmark(name, keyExpr, v, partitionNum) + } + } + } + } + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + if (mainArgs.contains("--nested-hash-only")) { + runNestedHashKeyBenchmarks() + return + } // nested type shuffle val numRows = 1000 @@ -567,24 +630,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase { } } - runBenchmarkWithTable("Nested hash partitioning key", 1024 * 1024 * 1) { v => - // Shapes whose leaves are primitives take the vectorized element path; the last two force - // the per-element path in the native kernel. - val shapes = Seq( - "struct" -> "named_struct('a', c1, 'b', CAST(c1 AS STRING))", - "array" -> "ARRAY_REPEAT(c1, 10)", - "struct, string>" -> - "named_struct('a', ARRAY_REPEAT(c1, 10), 'b', CAST(c1 AS STRING))", - "array>" -> - "ARRAY_REPEAT(named_struct('a', c1, 'b', CAST(c1 AS STRING)), 10)", - "struct, int>" -> - "named_struct('m', MAP(CAST(c1 AS STRING), c1), 'i', c1)") - shapes.foreach { case (name, keyExpr) => - Seq(5, manyPartitions).foreach { partitionNum => - shuffleNestedHashKeyBenchmark(name, keyExpr, v, partitionNum) - } - } - } + runNestedHashKeyBenchmarks() runBenchmarkWithTable("Shuffle on array", 1024 * 1024 * 1) { v => benchmarkTypes.foreach { dataType =>