comments and style

Signed-off-by: Haoyang Li <haoyangl@nvidia.com>

Author: thirtiseven

integration_tests/src/main/python/protobuf_test.py

@@ -1607,6 +1607,31 @@ def run_on_spark(spark):
     assert_gpu_and_cpu_error(run_on_spark, conf={}, error_message="Malformed")
 
 
+@pytest.mark.skipif(is_before_spark_340(), reason="from_protobuf is Spark 3.4.0+")
+@ignore_order(local=True)
+def test_from_protobuf_required_field_missing_permissive(spark_tmp_path, from_protobuf_fn):
+    """Required-field violations should null the whole row in PERMISSIVE mode."""
+    desc_path, desc_bytes = _setup_protobuf_desc(
+        spark_tmp_path, "required.desc", _build_required_field_descriptor_set_bytes)
+    message_name = "test.WithRequired"
+
+    missing_required_row = _encode_tag(2, 2) + _encode_varint(5) + b"hello"
+
+    def run_on_spark(spark):
+        df = spark.createDataFrame([(missing_required_row,)], schema="bin binary")
+        decoded = _call_from_protobuf(
+            from_protobuf_fn, f.col("bin"), message_name, desc_path, desc_bytes,
+            options={"mode": "PERMISSIVE"})
+        return df.select(
+            decoded.isNull().alias("decoded_is_null"),
+            decoded.getField("id").alias("id"),
+            decoded.getField("name").alias("name"),
+            decoded.getField("count").alias("count")
+        )
+
+    assert_gpu_and_cpu_are_equal_collect(run_on_spark)
+
+
 @pytest.mark.skipif(is_before_spark_340(), reason="from_protobuf is Spark 3.4.0+")
 @ignore_order(local=True)
 def test_from_protobuf_nested_required_field_missing_permissive(

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/protobuf/ProtobufSchemaValidator.scala

@@ -17,6 +17,7 @@
 package org.apache.spark.sql.rapids.protobuf
 
 import ai.rapids.cudf.DType
+
 import org.apache.spark.sql.rapids.GpuFromProtobuf
 import org.apache.spark.sql.types._
 

sql-plugin/src/main/spark340/scala/com/nvidia/spark/rapids/shims/SparkProtobufCompat.scala

@@ -230,13 +230,18 @@ private[shims] object SparkProtobufCompat extends Logging {
     override lazy val defaultValueResult: Either[String, Option[ProtobufDefaultValue]] =
       Try {
         if (PbReflect.hasDefaultValue(raw)) {
-          PbReflect.getDefaultValue(raw).map(toDefaultValue(_, protoTypeName, enumMetadata))
+          PbReflect.getDefaultValue(raw) match {
+            case Some(default) =>
+              toDefaultValue(default, protoTypeName, enumMetadata).map(Some(_))
+            case None =>
+              Right(None)
+          }
         } else {
-          None
+          Right(None)
         }
       }.toEither.left.map { t =>
         s"Failed to read protobuf default value for field '$name': ${t.getMessage}"
-      }
+      }.flatMap(identity)
     override lazy val messageDescriptor: Option[ProtobufMessageDescriptor] =
       if (protoTypeName == "MESSAGE") {
         Some(new ReflectiveMessageDescriptor(PbReflect.getMessageType(raw)))
@@ -248,26 +253,31 @@ private[shims] object SparkProtobufCompat extends Logging {
   private def toDefaultValue(
       rawDefault: AnyRef,
       protoTypeName: String,
-      enumMetadata: Option[ProtobufEnumMetadata]): ProtobufDefaultValue = protoTypeName match {
+      enumMetadata: Option[ProtobufEnumMetadata]): Either[String, ProtobufDefaultValue] =
+    protoTypeName match {
     case "BOOL" =>
-      ProtobufDefaultValue.BoolValue(rawDefault.asInstanceOf[java.lang.Boolean].booleanValue())
+      Right(ProtobufDefaultValue.BoolValue(
+        rawDefault.asInstanceOf[java.lang.Boolean].booleanValue()))
     case "FLOAT" =>
-      ProtobufDefaultValue.FloatValue(rawDefault.asInstanceOf[java.lang.Float].floatValue())
+      Right(ProtobufDefaultValue.FloatValue(
+        rawDefault.asInstanceOf[java.lang.Float].floatValue()))
     case "DOUBLE" =>
-      ProtobufDefaultValue.DoubleValue(rawDefault.asInstanceOf[java.lang.Double].doubleValue())
+      Right(ProtobufDefaultValue.DoubleValue(
+        rawDefault.asInstanceOf[java.lang.Double].doubleValue()))
     case "STRING" =>
-      ProtobufDefaultValue.StringValue(if (rawDefault == null) null else rawDefault.toString)
+      Right(ProtobufDefaultValue.StringValue(
+        if (rawDefault == null) null else rawDefault.toString))
     case "BYTES" =>
-      ProtobufDefaultValue.BinaryValue(extractBytes(rawDefault))
+      Right(ProtobufDefaultValue.BinaryValue(extractBytes(rawDefault)))
     case "ENUM" =>
       val number = extractNumber(rawDefault).intValue()
-      enumMetadata.map(_.enumDefault(number))
-        .getOrElse(ProtobufDefaultValue.EnumValue(number, rawDefault.toString))
+      Right(enumMetadata.map(_.enumDefault(number))
+        .getOrElse(ProtobufDefaultValue.EnumValue(number, rawDefault.toString)))
     case "INT32" | "UINT32" | "SINT32" | "FIXED32" | "SFIXED32" |
          "INT64" | "UINT64" | "SINT64" | "FIXED64" | "SFIXED64" =>
-      ProtobufDefaultValue.IntValue(extractNumber(rawDefault).longValue())
+      Right(ProtobufDefaultValue.IntValue(extractNumber(rawDefault).longValue()))
     case other =>
-      throw new IllegalStateException(
+      Left(
         s"Unsupported protobuf default value type '$other' for value ${rawDefault.toString}")
   }
 

sql-plugin/src/test/scala/com/nvidia/spark/rapids/shims/ProtobufExprShimsSuite.scala

@@ -125,7 +125,9 @@ class ProtobufExprShimsSuite extends AnyFunSuite {
   }
 
   private object FakeSpark35RetryFailureProtobufUtils {
-    def buildDescriptor(messageName: String, binaryFileDescriptorSet: Option[Array[Byte]]): String = {
+    def buildDescriptor(
+        messageName: String,
+        binaryFileDescriptorSet: Option[Array[Byte]]): String = {
       val bytes = binaryFileDescriptorSet.getOrElse(Array.emptyByteArray)
       if (bytes.sameElements(Array[Byte](1, 2, 3))) {
         throw new IllegalArgumentException(s"Unknown message $messageName")
@@ -280,6 +282,23 @@ class ProtobufExprShimsSuite extends AnyFunSuite {
     assert(SparkProtobufCompat.isGpuSupportedProtoSyntax("PROTO2"))
   }
 
+  test("compat returns Left for unsupported default value types") {
+    val method = SparkProtobufCompat.getClass.getDeclaredMethod(
+      "toDefaultValue",
+      classOf[AnyRef],
+      classOf[String],
+      classOf[scala.Option[_]])
+    method.setAccessible(true)
+
+    val result = method.invoke(
+      SparkProtobufCompat,
+      "opaque-default",
+      "MESSAGE",
+      scala.None).asInstanceOf[Either[String, ProtobufDefaultValue]]
+
+    assert(result.left.toOption.exists(_.contains("Unsupported protobuf default value type")))
+  }
+
   test("extractor preserves typed enum defaults") {
     val enumMeta = ProtobufEnumMetadata(Seq(
       ProtobufEnumValue(0, "UNKNOWN"),