feat: Add fuzz testing for arithmetic expressions

fuzz-testing/README.md

@@ -32,7 +32,6 @@ Planned areas of improvement:
 
 - ANSI mode
 - Support for all data types, expressions, and operators supported by Comet
-- Unary and binary arithmetic expressions
 - IF and CASE WHEN expressions
 - Complex (nested) expressions
 - Literal scalar values in queries

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

@@ -106,4 +106,8 @@ object Meta {
     Function("stddev_samp", 1),
     Function("corr", 2))
 
+  val unaryArithmeticOps: Seq[String] = Seq("+", "-")
+
+  val binaryArithmeticOps: Seq[String] = Seq("+", "-", "*", "/", "%", "&", "|", "^")
+
 }

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

@@ -42,11 +42,13 @@ object QueryGen {
     val uniqueQueries = mutable.HashSet[String]()
 
     for (_ <- 0 until numQueries) {
-      val sql = r.nextInt().abs % 4 match {
+      val sql = r.nextInt().abs % 6 match {
         case 0 => generateJoin(r, spark, numFiles)
         case 1 => generateAggregate(r, spark, numFiles)
         case 2 => generateScalar(r, spark, numFiles)
         case 3 => generateCast(r, spark, numFiles)
+        case 4 => generateUnaryArithmetic(r, spark, numFiles)
+        case 5 => generateBinaryArithmetic(r, spark, numFiles)
       }
       if (!uniqueQueries.contains(sql)) {
         uniqueQueries += sql
@@ -92,6 +94,33 @@ object QueryGen {
       s"ORDER BY ${args.mkString(", ")};"
   }
 
+  private def generateUnaryArithmetic(r: Random, spark: SparkSession, numFiles: Int): String = {
+    val tableName = s"test${r.nextInt(numFiles)}"
+    val table = spark.table(tableName)
+
+    val op = Utils.randomChoice(Meta.unaryArithmeticOps, r)
+    val a = Utils.randomChoice(table.columns, r)
+
+    // Example SELECT a, -a FROM test0
+    s"SELECT $a, $op$a " +
+      s"FROM $tableName " +
+      s"ORDER BY $a;"
+  }
+
+  private def generateBinaryArithmetic(r: Random, spark: SparkSession, numFiles: Int): String = {
+    val tableName = s"test${r.nextInt(numFiles)}"
+    val table = spark.table(tableName)
+
+    val op = Utils.randomChoice(Meta.binaryArithmeticOps, r)
+    val a = Utils.randomChoice(table.columns, r)
+    val b = Utils.randomChoice(table.columns, r)
+
+    // Example SELECT a, b, a+b FROM test0
+    s"SELECT $a, $b, $a $op $b " +
+      s"FROM $tableName " +
+      s"ORDER BY $a, $b;"
+  }
+
   private def generateCast(r: Random, spark: SparkSession, numFiles: Int): String = {
     val tableName = s"test${r.nextInt(numFiles)}"
     val table = spark.table(tableName)

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

@@ -19,7 +19,7 @@
 
 package org.apache.comet.fuzz
 
-import java.io.{BufferedWriter, FileWriter, PrintWriter}
+import java.io.{BufferedWriter, FileWriter, PrintWriter, StringWriter}
 
 import scala.io.Source
 
@@ -111,9 +111,11 @@ object QueryRunner {
                 showSQL(w, sql)
                 w.write(s"[ERROR] Query failed in Comet: ${e.getMessage}:\n")
                 w.write("```\n")
-                val p = new PrintWriter(w)
+                val sw = new StringWriter()
+                val p = new PrintWriter(sw)
                 e.printStackTrace(p)
                 p.close()
+                w.write(s"${sw.toString}\n")
                 w.write("```\n")
             }