sunchao commented on code in PR #6856:
URL: https://github.com/apache/datafusion-comet/pull/6856#discussion_r4238019814
##########
spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala:
##########
@@ -720,6 +727,63 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
}
}
+ /**
+ * Nested shapes the single-type array and struct groups above do not reach:
collections of
+ * strings and structs, maps, nulls at every level, and structs with
variable-width fields. Each
+ * shape is repartitioned by a round-robin so the shuffle moves the whole
column.
+ */
+ private def runNestedShapeBenchmarks(): Unit = {
+ val numRows = 1024 * 1024
+ val shapes: Seq[(String, String)] = Seq(
+ "array<string>" ->
+ "array(concat('a_', cast(id as string)), concat('b_', cast(id as
string)), 'c')",
+ "array<struct<int,string>>" ->
+ ("transform(sequence(0, cast(id % 4 as int)), i -> named_struct('i',
cast(id as int) + i, " +
+ "'s', concat('n_', cast(id as string))))"),
+ "array<array<int>>" ->
+ ("transform(sequence(0, cast(id % 3 as int)), i ->
transform(sequence(0, " +
+ "cast(id % 5 as int)), j -> cast(id as int) + i + j))"),
+ "array<int> (100 elements)" ->
+ "transform(sequence(1, 100), i -> cast(id % 1000 as int) + i)",
+ "map<string,int>" -> "map('k1', cast(id as int), 'k2', cast(id + 1 as
int))",
+ "map<int,array<string>>" ->
+ "map(cast(id % 10 as int), array(concat('v_', cast(id as string)),
'x'))",
+ "struct<int,string,string,long>" ->
+ ("named_struct('i', cast(id as int), 's1', concat('s1_', cast(id as
string)), " +
+ "'s2', concat('s2_', cast(id as string)), 'l', id)"),
+ "nullable struct, nullable fields" ->
+ ("if(id % 5 = 0, null, named_struct('i', if(id % 3 = 0, null, cast(id
as int)), " +
+ "'s', if(id % 4 = 0, null, concat('s_', cast(id as string)))))"),
+ "nullable array<int> with null elements" ->
+ ("if(id % 6 = 0, null, array(cast(id as int), if(id % 3 = 0, null,
cast(id + 1 as int)), " +
+ "cast(id + 2 as int)))"),
+ "struct<array<struct<long,string>>,string>" ->
+ ("named_struct('items', transform(sequence(0, cast(id % 3 as int)), i
-> " +
+ "named_struct('k', id + i, 'v', concat('v_', cast(i as string)))), "
+
+ "'label', concat('l_', cast(id as string)))"))
+
+ shapes.foreach { case (name, expr) =>
+ withTempPath { dir =>
+ val filename = dir.getCanonicalPath + "/nested.parquet"
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark
+ .range(numRows.toLong)
+ .selectExpr("id AS key", s"$expr AS c")
+ .write
+ .mode(SaveMode.Overwrite)
+ .parquet(filename)
+ }
+ Seq(5, manyPartitions).foreach { partitionNum =>
+ shuffleDeeplyNestedBenchmark(
Review Comment:
[P2] Make the added native-shuffle cases actually select native shuffle.
Each new shape calls `shuffleDeeplyNestedBenchmark`, whose `Comet (native
Shuffle)` arm uses `.repartition(partitionNum)` and sets
`spark.comet.shuffle.mode=native`. At both 5 and 201 partitions this requests
`RoundRobinPartitioning`, but
`spark.comet.shuffle.native.partitioning.roundrobin.enabled` remains false.
Consequently, the arm falls back to Spark shuffle instead of measuring native
shuffle, duplicating the Spark-shuffle comparison and leaving all 20 added
native measurements mislabeled. Please explicitly enable native round-robin, or
repartition these fixtures by their scalar `key`, and verify a
`CometNativeShuffle` exchange outside timing.
Evidence: Compiled the exact-head `CometConf.scala` and
`CometShuffleExchangeExec.scala` with a disposable Spark 4.1.3 planner probe.
For five representative added nested types at both partition counts, mode
`native` returned `None` from `shuffleSupported`, with reasons
`spark.comet.shuffle.native.partitioning.roundrobin.enabled is disabled` and
`Comet columnar shuffle not enabled`. Mode `jvm` returned
`Some(CometColumnarShuffle)`. Enabling the round-robin flag returned
`Some(CometNativeShuffle)`. Spark sources for every supported profile map
unkeyed repartitioning with more than one partition to
`RoundRobinPartitioning`. The benchmark and standard launcher do not enable
that flag.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
