null check for data generator

apache · Nov 5, 2022 · bcc333f · bcc333f
1 parent 66c6aab
commit bcc333f
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/...uf/src/test/scala/org/apache/spark/sql/protobuf/ProtobufCatalystDataConversionSuite.scala b/...uf/src/test/scala/org/apache/spark/sql/protobuf/ProtobufCatalystDataConversionSuite.scala
@@ -123,16 +123,21 @@ class ProtobufCatalystDataConversionSuite
     StringType -> ("StringMsg", ""))
 
   testingTypes.foreach { dt =>
-    val seed = 1 + scala.util.Random.nextInt((1024 - 1) + 1)
+    val seed = scala.util.Random.nextInt(RandomDataGenerator.MAX_STR_LEN)
     test(s"single $dt with seed $seed") {
 
       val (messageName, defaultValue) = catalystTypesToProtoMessages(dt.fields(0).dataType)
 
       val rand = new scala.util.Random(seed)
       val generator = RandomDataGenerator.forType(dt, rand = rand).get
       var data = generator()
-      while (data.asInstanceOf[Row].get(0) == defaultValue) // Do not use default values, since
-        data = generator()                                  // from_protobuf() returns null in v3.
+      // Do not use default values, since from_protobuf() returns null in v3.
+      while (
+        data != null &&
+        (data.asInstanceOf[Row].get(0) == defaultValue ||
+          (data.asInstanceOf[Row].get(0).isInstanceOf[Array[Byte]] &&
+          data.asInstanceOf[Row].get(0).asInstanceOf[Array[Byte]].isEmpty)))
+        data = generator()
 
       val converter = CatalystTypeConverters.createToCatalystConverter(dt)
       val input = Literal.create(converter(data), dt)