andygrove commented on code in PR #6856:
URL: https://github.com/apache/datafusion-comet/pull/6856#discussion_r4238504737
##########
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:
Fixed in 36acd72e8: the native arm turns on round-robin, and each Comet
arm's planned shuffle is checked before timing, skipping an arm that falls
back. All ten shapes now run native shuffle; the fuzz cases skip it because
Spark's scan reads that schema.
--
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]
