chore: Add shuffle benchmark for deeply nested schemas (apache#2902)

andygrove · web-flow · commit 60c0f1ed4d50 · 2025-12-18T19:29:12.000-07:00
diff --git a/dev/benchmarks/comet-tpch.sh b/dev/benchmarks/comet-tpch.sh
@@ -50,5 +50,4 @@ $SPARK_HOME/bin/spark-submit \
     --data $TPCH_DATA \
     --queries $TPCH_QUERIES \
     --output . \
-    --write /tmp \
     --iterations 1
diff --git a/spark/src/main/scala/org/apache/comet/testing/FuzzDataGenerator.scala b/spark/src/main/scala/org/apache/comet/testing/FuzzDataGenerator.scala
@@ -88,6 +88,70 @@ object FuzzDataGenerator {
     StructType(fields.toSeq)
   }
 
+  def generateNestedSchema(
+      r: Random,
+      numCols: Int,
+      minDepth: Int,
+      maxDepth: Int,
+      options: SchemaGenOptions): StructType = {
+    assert(numCols > 0)
+    assert(minDepth >= 0)
+    assert(maxDepth >= 0)
+    assert(minDepth <= maxDepth)
+    assert(
+      options.generateArray || options.generateStruct || options.generateMap,
+      "cannot generate nested schema if options do not include generating complex types")
+
+    var counter = 0
+
+    def generateFieldName() = {
+      val name = s"c_$counter"
+      counter += 1
+      name
+    }
+
+    def generateArray(depth: Int, name: String) = {
+      val element = genField(r, depth + 1)
+      StructField(name, DataTypes.createArrayType(element.dataType, true))
+    }
+
+    def generateStruct(depth: Int, name: String) = {
+      val fields =
+        Range(1, 2 + r.nextInt(10)).map(_ => genField(r, depth + 1)).toArray
+      StructField(name, DataTypes.createStructType(fields))
+    }
+
+    def generateMap(depth: Int, name: String) = {
+      val keyField = genField(r, depth + 1)
+      val valueField = genField(r, depth + 1)
+      StructField(name, DataTypes.createMapType(keyField.dataType, valueField.dataType))
+    }
+
+    def generatePrimitive(name: String) = {
+      StructField(name, randomChoice(options.primitiveTypes, r))
+    }
+
+    def genField(r: Random, depth: Int): StructField = {
+      val name = generateFieldName()
+      val generators = new ListBuffer[() => StructField]()
+      if (options.generateArray && depth < maxDepth) {
+        generators += (() => generateArray(depth + 1, name))
+      }
+      if (options.generateStruct && depth < maxDepth) {
+        generators += (() => generateStruct(depth + 1, name))
+      }
+      if (options.generateMap && depth < maxDepth) {
+        generators += (() => generateMap(depth + 1, name))
+      }
+      if (depth >= minDepth) {
+        generators += (() => generatePrimitive(name))
+      }
+      randomChoice(generators.toSeq, r)()
+    }
+
+    StructType(Range(0, numCols).map(_ => genField(r, 0)))
+  }
+
   def generateDataFrame(
       r: Random,
       spark: SparkSession,
diff --git a/spark/src/test/scala/org/apache/comet/DataGeneratorSuite.scala b/spark/src/test/scala/org/apache/comet/DataGeneratorSuite.scala
@@ -19,11 +19,46 @@
 
 package org.apache.comet
 
+import scala.util.Random
+
 import org.apache.spark.sql.CometTestBase
-import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructType}
+
+import org.apache.comet.testing.{FuzzDataGenerator, SchemaGenOptions}
 
 class DataGeneratorSuite extends CometTestBase {
 
+  test("generate nested schema has at least minDepth levels") {
+    val minDepth = 3
+    val numCols = 4
+    val schema = FuzzDataGenerator.generateNestedSchema(
+      new Random(42),
+      numCols,
+      minDepth = minDepth,
+      maxDepth = minDepth + 1,
+      options = SchemaGenOptions(generateMap = true, generateArray = true, generateStruct = true))
+    assert(schema.fields.length == numCols)
+
+    def calculateDepth(dataType: DataType): Int = {
+      dataType match {
+        case ArrayType(elementType, _) => 1 + calculateDepth(elementType)
+        case StructType(fields) =>
+          if (fields.isEmpty) 1
+          else 1 + fields.map(f => calculateDepth(f.dataType)).max
+        case MapType(k, v, _) =>
+          calculateDepth(k).max(calculateDepth(v))
+        case _ =>
+          // primitive type
+          1
+      }
+    }
+
+    val actualDepth = schema.fields.map(f => calculateDepth(f.dataType)).max
+    assert(
+      actualDepth >= minDepth,
+      s"Generated schema depth $actualDepth is less than required minimum depth $minDepth")
+  }
+
   test("test configurable stringGen in row generator") {
     val gen = DataGenerator.DEFAULT
     val chars = "abcde"
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
@@ -19,22 +19,29 @@
 
 package org.apache.spark.sql.benchmark
 
+import java.text.SimpleDateFormat
+
+import scala.util.Random
+
 import org.apache.spark.SparkConf
 import org.apache.spark.benchmark.Benchmark
-import org.apache.spark.sql.{Column, SparkSession}
+import org.apache.spark.sql.{Column, SaveMode, SparkSession}
 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.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` Results will be written to
- * "spark/benchmarks/CometShuffleBenchmark-**results.txt".
+ * `SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometShuffleBenchmark`
+ * Results will be written to "spark/benchmarks/CometShuffleBenchmark-**results.txt".
  */
+// spotless:on
 object CometShuffleBenchmark extends CometBenchmarkBase {
+
   override def getSparkSession: SparkSession = {
     val conf = new SparkConf()
       .setAppName("CometShuffleBenchmark")
@@ -97,7 +104,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
           }
         }
 
-        benchmark.addCase("SQL Parquet - Comet (Comet Arrow Shuffle)") { _ =>
+        benchmark.addCase("SQL Parquet - Comet (Comet JVM Shuffle)") { _ =>
           withSQLConf(
             CometConf.COMET_ENABLED.key -> "true",
             CometConf.COMET_EXEC_ENABLED.key -> "true",
@@ -154,7 +161,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
           }
         }
 
-        benchmark.addCase("SQL Parquet - Comet (Comet Arrow Shuffle)") { _ =>
+        benchmark.addCase("SQL Parquet - Comet (Comet JVM Shuffle)") { _ =>
           withSQLConf(
             CometConf.COMET_ENABLED.key -> "true",
             CometConf.COMET_EXEC_ENABLED.key -> "true",
@@ -209,7 +216,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
           }
         }
 
-        benchmark.addCase("SQL Parquet - Comet (Comet Arrow Shuffle)") { _ =>
+        benchmark.addCase("SQL Parquet - Comet (Comet JVM Shuffle)") { _ =>
           withSQLConf(
             CometConf.COMET_ENABLED.key -> "true",
             CometConf.COMET_EXEC_ENABLED.key -> "true",
@@ -224,7 +231,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
           }
         }
 
-        benchmark.addCase("SQL Parquet - Comet (Comet Arrow Shuffle + Prefer Dictionary)") { _ =>
+        benchmark.addCase("SQL Parquet - Comet (Comet JVM Shuffle + Prefer Dictionary)") { _ =>
           withSQLConf(
             CometConf.COMET_ENABLED.key -> "true",
             CometConf.COMET_EXEC_ENABLED.key -> "true",
@@ -239,7 +246,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
           }
         }
 
-        benchmark.addCase("SQL Parquet - Comet (Comet Arrow Shuffle + Fallback to string)") { _ =>
+        benchmark.addCase("SQL Parquet - Comet (Comet JVM Shuffle + Fallback to string)") { _ =>
           withSQLConf(
             CometConf.COMET_ENABLED.key -> "true",
             CometConf.COMET_EXEC_ENABLED.key -> "true",
@@ -305,7 +312,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
           }
         }
 
-        benchmark.addCase("SQL Parquet - Comet (Comet Arrow Shuffle)") { _ =>
+        benchmark.addCase("SQL Parquet - Comet (Comet JVM Shuffle)") { _ =>
           withSQLConf(
             CometConf.COMET_ENABLED.key -> "true",
             CometConf.COMET_EXEC_ENABLED.key -> "true",
@@ -319,7 +326,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
           }
         }
 
-        benchmark.addCase("SQL Parquet - Comet (Comet Async Arrow Shuffle)") { _ =>
+        benchmark.addCase("SQL Parquet - Comet (Comet Async JVM Shuffle)") { _ =>
           withSQLConf(
             CometConf.COMET_ENABLED.key -> "true",
             CometConf.COMET_EXEC_ENABLED.key -> "true",
@@ -389,7 +396,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
           }
         }
 
-        benchmark.addCase("SQL Parquet - Comet (Comet Arrow Shuffle)") { _ =>
+        benchmark.addCase("SQL Parquet - Comet (Comet JVM Shuffle)") { _ =>
           withSQLConf(
             CometConf.COMET_ENABLED.key -> "true",
             CometConf.COMET_EXEC_ENABLED.key -> "true",
@@ -402,7 +409,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
           }
         }
 
-        benchmark.addCase("SQL Parquet - Comet (Comet Shuffle)") { _ =>
+        benchmark.addCase("SQL Parquet - Comet (Comet Native Shuffle)") { _ =>
           withSQLConf(
             CometConf.COMET_ENABLED.key -> "true",
             CometConf.COMET_EXEC_ENABLED.key -> "true",
@@ -459,7 +466,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
           }
         }
 
-        benchmark.addCase("SQL Parquet - Comet (Comet Arrow Shuffle)") { _ =>
+        benchmark.addCase("SQL Parquet - Comet (Comet JVM Shuffle)") { _ =>
           withSQLConf(
             CometConf.COMET_ENABLED.key -> "true",
             CometConf.COMET_EXEC_ENABLED.key -> "true",
@@ -472,7 +479,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
           }
         }
 
-        benchmark.addCase("SQL Parquet - Comet (Comet Shuffle)") { _ =>
+        benchmark.addCase("SQL Parquet - Comet (Comet Native Shuffle)") { _ =>
           withSQLConf(
             CometConf.COMET_ENABLED.key -> "true",
             CometConf.COMET_EXEC_ENABLED.key -> "true",
@@ -490,7 +497,73 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
     }
   }
 
+  def shuffleDeeplyNestedBenchmark(
+      name: String,
+      filename: String,
+      numRows: Int,
+      partitionNum: Int): Unit = {
+    val benchmark =
+      new Benchmark(s"Shuffle with nested schema ($name)", numRows, output = output)
+    val df = spark.read.parquet(filename)
+    withTempTable("deeplyNestedTable") {
+      df.createOrReplaceTempView("deeplyNestedTable")
+      val sql = "select * from deeplyNestedTable"
+
+      benchmark.addCase("Spark") { _ =>
+        spark
+          .sql(sql)
+          .repartition(partitionNum)
+          .noop()
+      }
+
+      benchmark.addCase("Comet (Spark Shuffle)") { _ =>
+        withSQLConf(
+          CometConf.COMET_ENABLED.key -> "true",
+          CometConf.COMET_EXEC_ENABLED.key -> "true",
+          CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "false") {
+          spark
+            .sql(sql)
+            .repartition(partitionNum)
+            .noop()
+        }
+      }
+
+      for (shuffle <- Seq("jvm", "native")) {
+        benchmark.addCase(s"Comet ($shuffle Shuffle)") { _ =>
+          withSQLConf(
+            CometConf.COMET_ENABLED.key -> "true",
+            CometConf.COMET_EXEC_ENABLED.key -> "true",
+            CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
+            CometConf.COMET_SHUFFLE_MODE.key -> shuffle) {
+            spark
+              .sql(sql)
+              .repartition(partitionNum)
+              .noop()
+          }
+        }
+      }
+
+      benchmark.run()
+    }
+  }
+
   override def runCometBenchmark(mainArgs: Array[String]): Unit = {
+
+    // nested type shuffle
+    val numRows = 1000
+    for (maxDepth <- Seq(2, 6)) {
+      val filename =
+        createDeeplyNestedParquetFile(numRows, maxDepth)
+      try {
+        for (partitionNum <- Seq(5, 201)) {
+          val name = s"maxDepth=$maxDepth, partitionNum=$partitionNum"
+          shuffleDeeplyNestedBenchmark(name, filename, numRows, partitionNum)
+        }
+      } finally {
+        new java.io.File(filename).delete()
+      }
+    }
+
     runBenchmarkWithTable("Shuffle on array", 1024 * 1024 * 1) { v =>
       Seq(
         BooleanType,
@@ -725,4 +798,24 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
         }
     }
   }
+
+  private def createDeeplyNestedParquetFile(numRows: Int, maxDepth: Int): String = {
+    val r = new Random(42)
+    val options =
+      SchemaGenOptions(generateArray = true, generateStruct = true, generateMap = true)
+    val schema = FuzzDataGenerator.generateNestedSchema(r, 100, maxDepth - 1, maxDepth, options)
+    val tempDir = System.getProperty("java.io.tmpdir")
+    val filename = s"$tempDir/CometShuffleBenchmark_${System.currentTimeMillis()}.parquet"
+    withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+      val dataGenOptions = DataGenOptions(
+        generateNegativeZero = false,
+        // override base date due to known issues with experimental scans
+        baseDate =
+          new SimpleDateFormat("YYYY-MM-DD hh:mm:ss").parse("2024-05-25 12:34:56").getTime)
+      val df =
+        FuzzDataGenerator.generateDataFrame(r, spark, schema, numRows, dataGenOptions)
+      df.write.mode(SaveMode.Overwrite).parquet(filename)
+    }
+    filename
+  }
 }
