more refactoring

andygrove · May 29, 2024 · 0d69e10 · 0d69e10
1 parent e7a4a03
commit 0d69e10
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 10 deletions.
diff --git a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/DataGen.scala b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/DataGen.scala
@@ -19,13 +19,14 @@
 
 package org.apache.comet.fuzz
 
+import java.math.BigDecimal
 import java.nio.charset.Charset
 import java.sql.Timestamp
 
 import scala.util.Random
 
 import org.apache.spark.sql.{Row, SaveMode, SparkSession}
-import org.apache.spark.sql.types.{DataType, DataTypes, StructField, StructType}
+import org.apache.spark.sql.types.{DataType, DataTypes, DecimalType, StructField, StructType}
 
 object DataGen {
 
@@ -118,6 +119,8 @@ object DataGen {
             case _ => r.nextDouble()
           }
         })
+      case dt: DecimalType =>
+        Range(0, numRows).map(_ => new BigDecimal(r.nextDouble()).setScale(dt.scale))
       case DataTypes.StringType =>
         Range(0, numRows).map(_ => {
           r.nextInt(10) match {

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
@@ -32,16 +32,14 @@ object Meta {
     (DataTypes.LongType, 0.2),
     (DataTypes.FloatType, 0.2),
     (DataTypes.DoubleType, 0.2),
-    // TODO add Decimal support
-    // (DataTypes.createDecimalType(10,2), 0.2),
+    (DataTypes.createDecimalType(10, 2), 0.2),
     (DataTypes.DateType, 0.2),
     (DataTypes.TimestampType, 0.2),
     (DataTypes.TimestampNTZType, 0.2),
     (DataTypes.StringType, 0.2),
     (DataTypes.BinaryType, 0.2))
 
-  val scalarFunc: Seq[Function] = Seq(
-    // string
+  val stringScalarFunc: Seq[Function] = Seq(
     Function("substring", 3),
     Function("coalesce", 1),
     Function("starts_with", 2),
@@ -64,13 +62,17 @@ object Meta {
     Function("reverse", 1),
     Function("in_str", 2),
     Function("replace", 2),
-    Function("translate", 2),
-    // date
+    Function("translate", 2)
+  )
+
+  val dateScalarFunc: Seq[Function] = Seq(
     Function("year", 1),
     Function("hour", 1),
     Function("minute", 1),
-    Function("second", 1),
-    // math
+    Function("second", 1)
+  )
+
+  val mathScalarFunc: Seq[Function] = Seq(
     Function("abs", 1),
     Function("acos", 1),
     Function("asin", 1),
@@ -88,7 +90,10 @@ object Meta {
     Function("Sqrt", 1),
     Function("Tan", 1),
     Function("Ceil", 1),
-    Function("Floor", 1))
+    Function("Floor", 1)
+  )
+
+  val scalarFunc: Seq[Function] = stringScalarFunc ++ dateScalarFunc ++ mathScalarFunc
 
   val aggFunc: Seq[Function] = Seq(
     Function("min", 1),