chore: add TPC queries to be run by fuzzer correctness checker

comphead · comphead · commit 91f9ff10a477 · 2025-10-22T15:04:04.000-07:00
diff --git a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Main.scala b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Main.scala
@@ -61,6 +61,17 @@ class Conf(arguments: Seq[String]) extends ScallopConf(arguments) {
       opt[Int](required = false, descr = "Number of input files to use")
   }
   addSubcommand(runQueries)
+  object runTPCQueries extends Subcommand("runTPC") {
+    val dataFolder: ScallopOption[String] =
+      opt[String](
+        required = true,
+        descr = "Folder for input data. Expected folder struct `$dataFolder/tableName/*.parquet`")
+    val queriesFolder: ScallopOption[String] =
+      opt[String](
+        required = true,
+        descr = "Folder for test queries. Expected folder struct `$queriesFolder/*.sql`")
+  }
+  addSubcommand(runTPCQueries)
   verify()
 }
 
@@ -104,6 +115,11 @@ object Main {
           conf.generateQueries.numQueries())
       case Some(conf.runQueries) =>
         QueryRunner.runQueries(spark, conf.runQueries.numFiles(), conf.runQueries.filename())
+      case Some(conf.runTPCQueries) =>
+        QueryRunner.runTPCQueries(
+          spark,
+          conf.runTPCQueries.dataFolder(),
+          conf.runTPCQueries.queriesFolder())
       case _ =>
         // scalastyle:off println
         println("Invalid subcommand")
diff --git a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/QueryRunner.scala b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/QueryRunner.scala
@@ -19,27 +19,105 @@
 
 package org.apache.comet.fuzz
 
-import java.io.{BufferedWriter, FileWriter, PrintWriter, StringWriter}
+import java.io.{BufferedWriter, File, FileWriter, PrintWriter, StringWriter}
 
-import scala.collection.mutable.WrappedArray
+import scala.collection.mutable
 import scala.io.Source
 
 import org.apache.spark.sql.{Row, SparkSession}
 
 object QueryRunner {
 
+  def createOutputMdFile(): BufferedWriter = {
+    val outputFilename = s"results-${System.currentTimeMillis()}.md"
+    // scalastyle:off println
+    println(s"Writing results to $outputFilename")
+    // scalastyle:on println
+
+    new BufferedWriter(new FileWriter(outputFilename))
+  }
+
+  def assertCorrectness(
+      spark: SparkSession,
+      sql: String,
+      showFailedSparkQueries: Boolean = false,
+      output: BufferedWriter): Unit = {
+    try {
+      // execute with Spark
+      spark.conf.set("spark.comet.enabled", "false")
+      val df = spark.sql(sql)
+      val sparkRows = df.collect()
+      val sparkPlan = df.queryExecution.executedPlan.toString
+
+      // execute with Comet
+      try {
+        spark.conf.set("spark.comet.enabled", "true")
+        // complex type support until we support it natively
+        spark.conf.set("spark.comet.sparkToColumnar.enabled", "true")
+        spark.conf.set("spark.comet.convert.parquet.enabled", "true")
+        val df = spark.sql(sql)
+        val cometRows = df.collect()
+        val cometPlan = df.queryExecution.executedPlan.toString
+
+        if (sparkRows.length == cometRows.length) {
+          var i = 0
+          while (i < sparkRows.length) {
+            val l = sparkRows(i)
+            val r = cometRows(i)
+            assert(l.length == r.length)
+            for (j <- 0 until l.length) {
+              if (!same(l(j), r(j))) {
+                showSQL(output, sql)
+                showPlans(output, sparkPlan, cometPlan)
+                output.write(s"First difference at row $i:\n")
+                output.write("Spark: `" + formatRow(l) + "`\n")
+                output.write("Comet: `" + formatRow(r) + "`\n")
+                i = sparkRows.length
+              }
+            }
+            i += 1
+          }
+        } else {
+          showSQL(output, sql)
+          showPlans(output, sparkPlan, cometPlan)
+          output.write(
+            s"[ERROR] Spark produced ${sparkRows.length} rows and " +
+              s"Comet produced ${cometRows.length} rows.\n")
+        }
+      } catch {
+        case e: Exception =>
+          // the query worked in Spark but failed in Comet, so this is likely a bug in Comet
+          showSQL(output, sql)
+          output.write(s"[ERROR] Query failed in Comet: ${e.getMessage}:\n")
+          output.write("```\n")
+          val sw = new StringWriter()
+          val p = new PrintWriter(sw)
+          e.printStackTrace(p)
+          p.close()
+          output.write(s"${sw.toString}\n")
+          output.write("```\n")
+      }
+
+      // flush after every query so that results are saved in the event of the driver crashing
+      output.flush()
+
+    } catch {
+      case e: Exception =>
+        // we expect many generated queries to be invalid
+        if (showFailedSparkQueries) {
+          showSQL(output, sql)
+          output.write(s"Query failed in Spark: ${e.getMessage}\n")
+        }
+    }
+  }
+
   def runQueries(
       spark: SparkSession,
       numFiles: Int,
       filename: String,
       showFailedSparkQueries: Boolean = false): Unit = {
 
-    val outputFilename = s"results-${System.currentTimeMillis()}.md"
-    // scalastyle:off println
-    println(s"Writing results to $outputFilename")
-    // scalastyle:on println
-
-    val w = new BufferedWriter(new FileWriter(outputFilename))
+    val w = createOutputMdFile()
 
     // register input files
     for (i <- 0 until numFiles) {
@@ -55,104 +133,100 @@ object QueryRunner {
     try {
       querySource
         .getLines()
-        .foreach(sql => {
-
-          try {
-            // execute with Spark
-            spark.conf.set("spark.comet.enabled", "false")
-            val df = spark.sql(sql)
-            val sparkRows = df.collect()
-            val sparkPlan = df.queryExecution.executedPlan.toString
-
-            // execute with Comet
-            try {
-              spark.conf.set("spark.comet.enabled", "true")
-              // complex type support until we support it natively
-              spark.conf.set("spark.comet.sparkToColumnar.enabled", "true")
-              spark.conf.set("spark.comet.convert.parquet.enabled", "true")
-              val df = spark.sql(sql)
-              val cometRows = df.collect()
-              val cometPlan = df.queryExecution.executedPlan.toString
-
-              if (sparkRows.length == cometRows.length) {
-                var i = 0
-                while (i < sparkRows.length) {
-                  val l = sparkRows(i)
-                  val r = cometRows(i)
-                  assert(l.length == r.length)
-                  for (j <- 0 until l.length) {
-                    if (!same(l(j), r(j))) {
-                      showSQL(w, sql)
-                      showPlans(w, sparkPlan, cometPlan)
-                      w.write(s"First difference at row $i:\n")
-                      w.write("Spark: `" + formatRow(l) + "`\n")
-                      w.write("Comet: `" + formatRow(r) + "`\n")
-                      i = sparkRows.length
-                    }
-                  }
-                  i += 1
-                }
-              } else {
-                showSQL(w, sql)
-                showPlans(w, sparkPlan, cometPlan)
-                w.write(
-                  s"[ERROR] Spark produced ${sparkRows.length} rows and " +
-                    s"Comet produced ${cometRows.length} rows.\n")
-              }
-            } catch {
-              case e: Exception =>
-                // the query worked in Spark but failed in Comet, so this is likely a bug in Comet
-                showSQL(w, sql)
-                w.write(s"[ERROR] Query failed in Comet: ${e.getMessage}:\n")
-                w.write("```\n")
-                val sw = new StringWriter()
-                val p = new PrintWriter(sw)
-                e.printStackTrace(p)
-                p.close()
-                w.write(s"${sw.toString}\n")
-                w.write("```\n")
-            }
-
-            // flush after every query so that results are saved in the event of the driver crashing
-            w.flush()
-
-          } catch {
-            case e: Exception =>
-              // we expect many generated queries to be invalid
-              if (showFailedSparkQueries) {
-                showSQL(w, sql)
-                w.write(s"Query failed in Spark: ${e.getMessage}\n")
-              }
-          }
-        })
+        .foreach(sql => assertCorrectness(spark, sql, showFailedSparkQueries, output = w))
 
     } finally {
       w.close()
       querySource.close()
     }
   }
 
+  def runTPCQueries(
+      spark: SparkSession,
+      dataFolderName: String,
+      queriesFolderName: String): Unit = {
+    val output = QueryRunner.createOutputMdFile()
+
+    // Load data tables from dataFolder
+    val dataFolder = new File(dataFolderName)
+    if (!dataFolder.exists() || !dataFolder.isDirectory) {
+      // scalastyle:off println
+      println(s"Error: Data folder $dataFolder does not exist or is not a directory")
+      // scalastyle:on println
+      sys.exit(-1)
+    }
+
+    // Traverse data folder and create temp views
+    dataFolder.listFiles().filter(_.isDirectory).foreach { tableDir =>
+      val tableName = tableDir.getName
+      val parquetPath = s"${tableDir.getAbsolutePath}/*.parquet"
+      spark.read.parquet(parquetPath).createOrReplaceTempView(tableName)
+      // scalastyle:off println
+      println(s"Created temp view: $tableName from $parquetPath")
+    // scalastyle:on println
+    }
+
+    // Load and run queries from queriesFolder
+    val queriesFolder = new File(queriesFolderName)
+    if (!queriesFolder.exists() || !queriesFolder.isDirectory) {
+      // scalastyle:off println
+      println(s"Error: Queries folder $queriesFolder does not exist or is not a directory")
+      // scalastyle:on println
+      sys.exit(-1)
+    }
+
+    // Traverse queries folder and run each .sql file
+    queriesFolder.listFiles().filter(f => f.isFile && f.getName.endsWith(".sql")).foreach {
+      sqlFile =>
+        // scalastyle:off println
+        println(s"Running query from: ${sqlFile.getName}")
+        // scalastyle:on println
+
+        val querySource = Source.fromFile(sqlFile)
+        try {
+          val sql = querySource.mkString
+          QueryRunner.assertCorrectness(spark, sql, showFailedSparkQueries = false, output)
+        } finally {
+          querySource.close()
+        }
+    }
+
+    output.close()
+  }
+
   private def same(l: Any, r: Any): Boolean = {
+    if (l == null || r == null) {
+      return l == null && r == null
+    }
     (l, r) match {
+      case (a: Float, b: Float) if a.isPosInfinity => b.isPosInfinity
+      case (a: Float, b: Float) if a.isNegInfinity => b.isNegInfinity
       case (a: Float, b: Float) if a.isInfinity => b.isInfinity
       case (a: Float, b: Float) if a.isNaN => b.isNaN
       case (a: Float, b: Float) => (a - b).abs <= 0.000001f
+      case (a: Double, b: Double) if a.isPosInfinity => b.isPosInfinity
+      case (a: Double, b: Double) if a.isNegInfinity => b.isNegInfinity
       case (a: Double, b: Double) if a.isInfinity => b.isInfinity
       case (a: Double, b: Double) if a.isNaN => b.isNaN
       case (a: Double, b: Double) => (a - b).abs <= 0.000001
       case (a: Array[_], b: Array[_]) =>
         a.length == b.length && a.zip(b).forall(x => same(x._1, x._2))
-      case (a: WrappedArray[_], b: WrappedArray[_]) =>
+      case (a: mutable.WrappedArray[_], b: mutable.WrappedArray[_]) =>
         a.length == b.length && a.zip(b).forall(x => same(x._1, x._2))
+      case (a: Row, b: Row) =>
+        val aa = a.toSeq
+        val bb = b.toSeq
+        aa.length == bb.length && aa.zip(bb).forall(x => same(x._1, x._2))
       case (a, b) => a == b
     }
   }
 
   private def format(value: Any): String = {
     value match {
       case null => "NULL"
-      case v: WrappedArray[_] => s"[${v.map(format).mkString(",")}]"
+      case v: mutable.WrappedArray[_] => s"[${v.map(format).mkString(",")}]"
       case v: Array[Byte] => s"[${v.mkString(",")}]"
+      case r: Row => formatRow(r)
       case other => other.toString
     }
   }
