chore: extract comparison tool from fuzzer

Author: comphead

fuzz-testing/src/main/scala/org/apache/comet/fuzz/QueryRunner.scala

@@ -21,13 +21,22 @@ package org.apache.comet.fuzz
 
 import java.io.{BufferedWriter, 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 runQueries(
       spark: SparkSession,
       numFiles: Int,
@@ -39,12 +48,7 @@ object QueryRunner {
     var cometFailureCount = 0
     var cometSuccessCount = 0
 
-    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) {
@@ -76,42 +80,13 @@ object QueryRunner {
               val cometRows = df.collect()
               val cometPlan = df.queryExecution.executedPlan.toString
 
-              var success = true
-              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))) {
-                      success = false
-                      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 {
-                success = false
-                showSQL(w, sql)
-                showPlans(w, sparkPlan, cometPlan)
-                w.write(
-                  s"[ERROR] Spark produced ${sparkRows.length} rows and " +
-                    s"Comet produced ${cometRows.length} rows.\n")
-              }
-
-              // check that the plan contains Comet operators
-              if (!cometPlan.contains("Comet")) {
-                success = false
-                showSQL(w, sql)
-                showPlans(w, sparkPlan, cometPlan)
-                w.write("[ERROR] Comet did not accelerate any part of the plan\n")
-              }
+              val success = QueryComparison.assertSameRows(
+                sparkRows,
+                cometRows,
+                sqlText = sql,
+                sparkPlan,
+                cometPlan,
+                output = w)
 
               if (success) {
                 cometSuccessCount += 1
@@ -123,7 +98,7 @@ object QueryRunner {
               case e: Exception =>
                 // the query worked in Spark but failed in Comet, so this is likely a bug in Comet
                 cometFailureCount += 1
-                showSQL(w, sql)
+                QueryComparison.showSQL(w, sql)
                 w.write("### Spark Plan\n")
                 w.write(s"```\n$sparkPlan\n```\n")
 
@@ -145,7 +120,7 @@ object QueryRunner {
               // we expect many generated queries to be invalid
               invalidQueryCount += 1
               if (showFailedSparkQueries) {
-                showSQL(w, sql)
+                QueryComparison.showSQL(w, sql)
                 w.write(s"Query failed in Spark: ${e.getMessage}\n")
               }
           }
@@ -161,6 +136,56 @@ object QueryRunner {
       querySource.close()
     }
   }
+}
+
+object QueryComparison {
+  def assertSameRows(
+      sparkRows: Array[Row],
+      cometRows: Array[Row],
+      sqlText: String,
+      sparkPlan: String,
+      cometPlan: String,
+      output: BufferedWriter): Boolean = {
+    var success = true
+    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))) {
+            success = false
+            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 {
+      success = false
+      showSQL(output, sqlText)
+      showPlans(output, sparkPlan, cometPlan)
+      output.write(
+        s"[ERROR] Spark produced ${sparkRows.length} rows and " +
+          s"Comet produced ${cometRows.length} rows.\n")
+    }
+
+    // check that the plan contains Comet operators
+    if (!cometPlan.contains("Comet")) {
+      success = false
+      showSQL(output, sqlText)
+      showPlans(output, sparkPlan, cometPlan)
+      output.write("[ERROR] Comet did not accelerate any part of the plan\n")
+    }
+
+    success
+
+  }
 
   private def same(l: Any, r: Any): Boolean = {
     if (l == null || r == null) {
@@ -179,7 +204,7 @@ object QueryRunner {
       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
@@ -192,7 +217,7 @@ object QueryRunner {
   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
@@ -203,7 +228,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(" ")
@@ -229,5 +254,4 @@ object QueryRunner {
     w.write("### Comet Plan\n")
     w.write(s"```\n$cometPlan\n```\n")
   }
-
 }