chore: extract comparison tool from fuzzer

comphead · comphead · commit d579011f44fd · 2025-10-23T11:04:26.000-07:00
diff --git a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/ComparisonTool.scala b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/ComparisonTool.scala
@@ -0,0 +1,157 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet.fuzz
+
+import java.io.File
+
+import scala.util.Random
+
+import org.rogach.scallop.{ScallopConf, Subcommand}
+import org.rogach.scallop.ScallopOption
+
+import org.apache.spark.sql.SparkSession
+
+import org.apache.comet.testing.{DataGenOptions, ParquetGenerator}
+
+class Conf(arguments: Seq[String]) extends ScallopConf(arguments) {
+  object compareParquet extends Subcommand("compareParquet") {
+    val inputSparkFolder: ScallopOption[String] =
+      opt[String](required = true, descr = "Folder with Spark produced results in Parquet format")
+    val inputCometFolder: ScallopOption[String] =
+      opt[String](required = true, descr = "Folder with Comet produced results in Parquet format")
+  }
+  addSubcommand(compareParquet)
+  verify()
+}
+
+object Main {
+
+  lazy val spark: SparkSession = SparkSession
+    .builder()
+    .getOrCreate()
+
+  def main(args: Array[String]): Unit = {
+    val conf = new Conf(args.toIndexedSeq)
+    conf.subcommand match {
+      case Some(conf.compareParquet) =>
+        compareParquetFolders(
+          spark,
+          conf.compareParquet.inputSparkFolder(),
+          conf.compareParquet.inputCometFolder())
+
+      case _ =>
+        // scalastyle:off println
+        println("Invalid subcommand")
+        // scalastyle:on println
+        sys.exit(-1)
+    }
+  }
+
+  private def compareParquetFolders(
+      spark: SparkSession,
+      sparkFolderPath: String,
+      cometFolderPath: String): Unit = {
+
+    val output = QueryRunner.createOutputMdFile()
+
+    try {
+      val sparkFolder = new File(sparkFolderPath)
+      val cometFolder = new File(cometFolderPath)
+
+      if (!sparkFolder.exists() || !sparkFolder.isDirectory) {
+        throw new IllegalArgumentException(
+          s"Spark folder does not exist or is not a directory: $sparkFolderPath")
+      }
+
+      if (!cometFolder.exists() || !cometFolder.isDirectory) {
+        throw new IllegalArgumentException(
+          s"Comet folder does not exist or is not a directory: $cometFolderPath")
+      }
+
+      // Get all subdirectories from the Spark folder
+      val sparkSubfolders = sparkFolder
+        .listFiles()
+        .filter(_.isDirectory)
+        .map(_.getName)
+        .sorted
+
+      output.write(s"# Comparing Parquet Folders\n\n")
+      output.write(s"Spark folder: $sparkFolderPath\n")
+      output.write(s"Comet folder: $cometFolderPath\n")
+      output.write(s"Found ${sparkSubfolders.length} subfolders to compare\n\n")
+
+      // Compare each subfolder
+      sparkSubfolders.foreach { subfolderName =>
+        val sparkSubfolderPath = new File(sparkFolder, subfolderName)
+        val cometSubfolderPath = new File(cometFolder, subfolderName)
+
+        if (!cometSubfolderPath.exists() || !cometSubfolderPath.isDirectory) {
+          output.write(s"## Subfolder: $subfolderName\n")
+          output.write(
+            s"[WARNING] Comet subfolder not found: ${cometSubfolderPath.getAbsolutePath}\n\n")
+        } else {
+          output.write(s"## Comparing subfolder: $subfolderName\n\n")
+
+          try {
+            // Read Spark parquet files
+            spark.conf.set("spark.comet.enabled", "false")
+            val sparkDf = spark.read.parquet(sparkSubfolderPath.getAbsolutePath)
+            val sparkRows = sparkDf.collect()
+            val sparkPlan = sparkDf.queryExecution.executedPlan.toString
+
+            // Read Comet parquet files
+            val cometDf = spark.read.parquet(cometSubfolderPath.getAbsolutePath)
+            val cometRows = cometDf.collect()
+            val cometPlan = cometDf.queryExecution.executedPlan.toString
+
+            // Compare the results
+            QueryComparison.assertSameRows(
+              sparkRows,
+              cometRows,
+              sqlText = s"Reading parquet from subfolder: $subfolderName",
+              sparkPlan,
+              cometPlan,
+              output)
+
+            output.write(s"✓ Subfolder $subfolderName: ${sparkRows.length} rows matched\n\n")
+
+          } catch {
+            case e: Exception =>
+              output.write(
+                s"[ERROR] Failed to compare subfolder $subfolderName: ${e.getMessage}\n")
+              val sw = new java.io.StringWriter()
+              val p = new java.io.PrintWriter(sw)
+              e.printStackTrace(p)
+              p.close()
+              output.write(s"```\n${sw.toString}\n```\n\n")
+          }
+        }
+
+        output.flush()
+      }
+
+      output.write(s"\n# Comparison Complete\n")
+      output.write(s"Compared ${sparkSubfolders.length} subfolders\n")
+
+    } finally {
+      output.close()
+    }
+  }
+}
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,17 +61,6 @@ 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()
 }
 
@@ -115,11 +104,6 @@ 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,7 +19,7 @@
 
 package org.apache.comet.fuzz
 
-import java.io.{BufferedWriter, File, FileWriter, PrintWriter, StringWriter}
+import java.io.{BufferedWriter, FileWriter, PrintWriter, StringWriter}
 
 import scala.collection.mutable
 import scala.io.Source
@@ -37,7 +37,7 @@ object QueryRunner {
     new BufferedWriter(new FileWriter(outputFilename))
   }
 
-  def assertCorrectness(
+  def executeSQLAndAssertCorrectness(
       spark: SparkSession,
       sql: String,
       showFailedSparkQueries: Boolean = false,
@@ -59,35 +59,17 @@ object QueryRunner {
         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")
-        }
+        QueryComparison.assertSameRows(
+          sparkRows,
+          cometRows,
+          sqlText = sql,
+          sparkPlan,
+          cometPlan,
+          output)
       } catch {
         case e: Exception =>
           // the query worked in Spark but failed in Comet, so this is likely a bug in Comet
-          showSQL(output, sql)
+          QueryComparison.showSQL(output, sql)
           output.write(s"[ERROR] Query failed in Comet: ${e.getMessage}:\n")
           output.write("```\n")
           val sw = new StringWriter()
@@ -105,7 +87,7 @@ object QueryRunner {
       case e: Exception =>
         // we expect many generated queries to be invalid
         if (showFailedSparkQueries) {
-          showSQL(output, sql)
+          QueryComparison.showSQL(output, sql)
           output.write(s"Query failed in Spark: ${e.getMessage}\n")
         }
     }
@@ -133,67 +115,18 @@ object QueryRunner {
     try {
       querySource
         .getLines()
-        .foreach(sql => assertCorrectness(spark, sql, showFailedSparkQueries, output = w))
+        .foreach(sql =>
+          executeSQLAndAssertCorrectness(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()
-  }
+}
 
+object QueryComparison {
   private def same(l: Any, r: Any): Boolean = {
     if (l == null || r == null) {
       return l == null && r == null
@@ -235,7 +168,7 @@ object QueryRunner {
     row.toSeq.map(format).mkString(",")
   }
 
-  private def showSQL(w: BufferedWriter, sql: String, maxLength: Int = 120): Unit = {
+  def showSQL(w: BufferedWriter, sql: String, maxLength: Int = 120): Unit = {
     w.write("## SQL\n")
     w.write("```\n")
     val words = sql.split(" ")
@@ -262,4 +195,37 @@ object QueryRunner {
     w.write(s"```\n$cometPlan\n```\n")
   }
 
+  def assertSameRows(
+      sparkRows: Array[Row],
+      cometRows: Array[Row],
+      sqlText: String,
+      sparkPlan: String,
+      cometPlan: String,
+      output: BufferedWriter): Unit = {
+    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, sqlText)
+            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, sqlText)
+      showPlans(output, sparkPlan, cometPlan)
+      output.write(
+        s"[ERROR] Spark produced ${sparkRows.length} rows and " +
+          s"Comet produced ${cometRows.length} rows.\n")
+    }
+  }
 }
