address comments

apache · May 23, 2024 · 14ab3ce · 14ab3ce
1 parent 34acfd2
commit 14ab3ce
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 12 deletions.
diff --git a/spark/src/test/scala/org/apache/comet/DataGenerator.scala b/spark/src/test/scala/org/apache/comet/DataGenerator.scala
@@ -19,7 +19,6 @@
 
 package org.apache.comet
 
-import scala.collection.mutable
 import scala.util.Random
 
 import org.apache.spark.sql.{RandomDataGenerator, Row}
@@ -106,27 +105,28 @@ class DataGenerator(r: Random) {
   // configured to generate strings by passing a stringGen function. Other types are delegated
   // to Spark's RandomDataGenerator.
   def generateRow(schema: StructType, stringGen: Option[() => String] = None): Row = {
-    val fields = mutable.ArrayBuffer.empty[Any]
-    schema.fields.foreach { f =>
+    val fields = schema.fields.map { f =>
       f.dataType match {
         case StructType(children) =>
-          fields += generateRow(StructType(children), stringGen)
+          generateRow(StructType(children), stringGen)
         case StringType if stringGen.isDefined =>
           val gen = stringGen.get
           val data = if (f.nullable && r.nextFloat() <= PROBABILITY_OF_NULL) {
             null
           } else {
             gen()
           }
-          fields += data
+          data
         case _ =>
-          val generator = RandomDataGenerator.forType(f.dataType, f.nullable, r)
-          assert(generator.isDefined, "Unsupported type")
-          val gen = generator.get
-          fields += gen()
+          val gen = RandomDataGenerator.forType(f.dataType, f.nullable, r) match {
+            case Some(g) => g
+            case None =>
+              throw new IllegalStateException(s"No RandomDataGenerator for type ${f.dataType}")
+          }
+          gen()
       }
-    }
-    Row.fromSeq(fields.toSeq)
+    }.toSeq
+    Row.fromSeq(fields)
   }
 
   def generateRows(

diff --git a/spark/src/test/scala/org/apache/comet/DataGeneratorSuite.scala b/spark/src/test/scala/org/apache/comet/DataGeneratorSuite.scala
@@ -43,7 +43,7 @@ class DataGeneratorSuite extends CometTestBase {
         }
       })
     // 0.1 null probability
-    assert(numNulls >= 5 && numNulls <= 15)
+    assert(numNulls >= 0.05 * numRows && numNulls <= 0.15 * numRows)
   }
 
 }