add unary math

apache · Jun 5, 2024 · 449cb92 · 449cb92
1 parent cfe8c72
commit 449cb92
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 6 deletions.
diff --git a/fuzz-testing/README.md b/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

diff --git a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Meta.scala b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Meta.scala
@@ -106,6 +106,8 @@ object Meta {
     Function("stddev_samp", 1),
     Function("corr", 2))
 
-  val arithmeticOps: Seq[String] = Seq("+", "-", "*", "/", "%")
+  val unaryArithmeticOps: Seq[String] = Seq("+", "-")
+
+  val binaryArithmeticOps: Seq[String] = Seq("+", "-", "*", "/", "%")
 
 }
diff --git a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/QueryGen.scala b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/QueryGen.scala
@@ -42,12 +42,13 @@ object QueryGen {
     val uniqueQueries = mutable.HashSet[String]()
 
     for (_ <- 0 until numQueries) {
-      val sql = r.nextInt().abs % 5 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 => generateArithmetic(r, spark, numFiles)
+        case 4 => generateUnaryArithmetic(r, spark, numFiles)
+        case 5 => generateBinaryArithmetic(r, spark, numFiles)
       }
       if (!uniqueQueries.contains(sql)) {
         uniqueQueries += sql
@@ -93,11 +94,24 @@ object QueryGen {
       s"ORDER BY ${args.mkString(", ")};"
   }
 
-  private def generateArithmetic(r: Random, spark: SparkSession, numFiles: Int): String = {
+  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.arithmeticOps, r)
+    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)