This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new fd8e09e699 bench: measure nested types as native shuffle hash
partitioning keys (#5788)
fd8e09e699 is described below
commit fd8e09e699e4e5e8a399616083d292dfb49cbb78
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Thu Sep 10 04:58:13 2026 +0800
bench: measure nested types as native shuffle hash partitioning keys (#5788)
* 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<int, string> 5 partitions 79ms -> 41ms 1.9x
struct<int, string> 201 partitions 126ms -> 89ms 1.4x
array<int> 5 partitions 76ms -> 52ms 1.5x
array<int> 201 partitions 145ms -> 101ms 1.4x
struct<array<int>, string> 5 partitions 100ms -> 68ms 1.5x
array<struct<int, string>> 5 partitions 164ms -> 561ms 0.3x
array<struct<int, string>> 201 partitions 250ms -> 647ms 0.4x
struct<map<string, int>, int> 5 partitions 121ms -> 97ms 1.2x
struct<map<string, int>, int> 201 partitions 160ms -> 155ms 1.0x
The native case is not silently falling back to Spark's shuffle. Rerunning
the
`array<struct<..>>` 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<struct<..>>` 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 <[email protected]>
* 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.
---------
Co-authored-by: Claude Code <[email protected]>
---
.../sql/benchmark/CometShuffleBenchmark.scala | 140 +++++++++++++++++++++
1 file changed, 140 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 9f5aaef306..64d9d0237a 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
@@ -475,7 +478,142 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
}
}
+ /**
+ * Nested hash partitioning keys, which native shuffle admits only when
+ * `spark.comet.shuffle.native.partitioning.hash.nested.enabled` is on.
+ *
+ * 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.
+ *
+ * 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,
+ 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()
+ }
+ }
+
+ 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")
+ // 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()
+ }
+ }
+ }
+
+ benchmark.run()
+ }
+ }
+ }
+
+ private def runNestedHashKeyBenchmarks(): Unit = {
+ runBenchmarkWithTable("Nested hash partitioning key", 1024 * 1024 * 1) { v
=>
+ val shapes = Seq(
+ "struct<int, string>" -> "named_struct('a', c1, 'b', CAST(c1 AS
STRING))",
+ "array<int>" -> "ARRAY_REPEAT(c1, 10)",
+ "struct<array<int>, string>" ->
+ "named_struct('a', ARRAY_REPEAT(c1, 10), 'b', CAST(c1 AS STRING))",
+ "array<struct<int, string>>" ->
+ "ARRAY_REPEAT(named_struct('a', c1, 'b', CAST(c1 AS STRING)), 10)",
+ "struct<map<string, int>, 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<map<string, int>, 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
@@ -492,6 +630,8 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
}
}
+ runNestedHashKeyBenchmarks()
+
runBenchmarkWithTable("Shuffle on array", 1024 * 1024 * 1) { v =>
benchmarkTypes.foreach { dataType =>
Seq(5, manyPartitions).foreach { partitionNum =>
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
