[SPARK-41396][SQL][PROTOBUF] OneOf field support and recursion checks

connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/ProtobufDataToCatalyst.scala

@@ -39,7 +39,7 @@ private[protobuf] case class ProtobufDataToCatalyst(
   override def inputTypes: Seq[AbstractDataType] = Seq(BinaryType)
 
   override lazy val dataType: DataType = {
-    val dt = SchemaConverters.toSqlType(messageDescriptor).dataType
+    val dt = SchemaConverters.toSqlType(messageDescriptor, protobufOptions).dataType
     parseMode match {
       // With PermissiveMode, the output Catalyst row might contain columns of null values for
       // corrupt records, even if some of the columns are not nullable in the user-provided schema.

connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/ProtobufDeserializer.scala

@@ -156,6 +156,9 @@ private[sql] class ProtobufDeserializer(
     (protoType.getJavaType, catalystType) match {
 
       case (null, NullType) => (updater, ordinal, _) => updater.setNullAt(ordinal)
+      // It is possible that this will result in data being dropped, This is intentional,
+      // to catch recursive fields and drop them as necessary.
+      case (MESSAGE, NullType) => (updater, ordinal, _) => updater.setNullAt(ordinal)
 
       // TODO: we can avoid boxing if future version of Protobuf provide primitive accessors.
       case (BOOLEAN, BooleanType) =>
@@ -171,7 +174,7 @@ private[sql] class ProtobufDeserializer(
         (updater, ordinal, value) => updater.setShort(ordinal, value.asInstanceOf[Short])
 
       case  (
-        BOOLEAN | INT | FLOAT | DOUBLE | LONG | STRING | ENUM | BYTE_STRING,
+        MESSAGE | BOOLEAN | INT | FLOAT | DOUBLE | LONG | STRING | ENUM | BYTE_STRING,
         ArrayType(dataType: DataType, containsNull)) if protoType.isRepeated =>
         newArrayWriter(protoType, protoPath, catalystPath, dataType, containsNull)
 
@@ -235,9 +238,6 @@ private[sql] class ProtobufDeserializer(
           writeRecord(new RowUpdater(row), value.asInstanceOf[DynamicMessage])
           updater.set(ordinal, row)
 
-      case (MESSAGE, ArrayType(st: StructType, containsNull)) =>
-        newArrayWriter(protoType, protoPath, catalystPath, st, containsNull)
-
       case (ENUM, StringType) =>
         (updater, ordinal, value) => updater.set(ordinal, UTF8String.fromString(value.toString))
 

connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/utils/ProtobufOptions.scala

@@ -38,6 +38,14 @@ private[sql] class ProtobufOptions(
 
   val parseMode: ParseMode =
     parameters.get("mode").map(ParseMode.fromString).getOrElse(FailFastMode)
+
+  // Setting the `recursive.fields.max.depth` to 0 drops all recursive fields,
+  // 1 allows it to be recurse once, and 2 allows it to be recursed twice and so on.
+  // A value of `recursive.fields.max.depth` greater than 10 is not permitted. If it is not
+  // specified, the default value is -1; recursive fields are not permitted. If a protobuf
+  // record has more depth than the allowed value for recursive fields, it will be truncated
+  // and some fields may be discarded.
+  val recursiveFieldMaxDepth: Int = parameters.getOrElse("recursive.fields.max.depth", "-1").toInt
 }
 
 private[sql] object ProtobufOptions {

connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/utils/SchemaConverters.scala

@@ -40,19 +40,27 @@ object SchemaConverters {
    *
    * @since 3.4.0
    */
-  def toSqlType(descriptor: Descriptor): SchemaType = {
-    toSqlTypeHelper(descriptor)
+  def toSqlType(
+      descriptor: Descriptor,
+      protobufOptions: ProtobufOptions = ProtobufOptions(Map.empty)): SchemaType = {
+    toSqlTypeHelper(descriptor, protobufOptions)
   }
 
-  def toSqlTypeHelper(descriptor: Descriptor): SchemaType = ScalaReflectionLock.synchronized {
+  def toSqlTypeHelper(
+      descriptor: Descriptor,
+      protobufOptions: ProtobufOptions): SchemaType = ScalaReflectionLock.synchronized {
     SchemaType(
-      StructType(descriptor.getFields.asScala.flatMap(structFieldFor(_, Set.empty)).toArray),
+      StructType(descriptor.getFields.asScala.flatMap(
+        structFieldFor(_,
+          Map(descriptor.getFullName -> 1),
+          protobufOptions: ProtobufOptions)).toArray),
       nullable = true)
   }
 
   def structFieldFor(
       fd: FieldDescriptor,
-      existingRecordNames: Set[String]): Option[StructField] = {
+      existingRecordNames: Map[String, Int],
+      protobufOptions: ProtobufOptions): Option[StructField] = {
     import com.google.protobuf.Descriptors.FieldDescriptor.JavaType._
     val dataType = fd.getJavaType match {
       case INT => Some(IntegerType)
@@ -81,9 +89,17 @@ object SchemaConverters {
         fd.getMessageType.getFields.forEach { field =>
           field.getName match {
             case "key" =>
-              keyType = structFieldFor(field, existingRecordNames).get.dataType
+              keyType =
+                structFieldFor(
+                  field,
+                  existingRecordNames,
+                  protobufOptions).get.dataType
             case "value" =>
-              valueType = structFieldFor(field, existingRecordNames).get.dataType
+              valueType =
+                structFieldFor(
+                  field,
+                  existingRecordNames,
+                  protobufOptions).get.dataType
           }
         }
         return Option(
@@ -92,17 +108,35 @@ object SchemaConverters {
             MapType(keyType, valueType, valueContainsNull = false).defaultConcreteType,
             nullable = false))
       case MESSAGE =>
-        if (existingRecordNames.contains(fd.getFullName)) {
+        // If the `recursive.fields.max.depth` value is not specified, it will default to -1;
+        // recursive fields are not permitted. Setting it to 0 drops all recursive fields,
+        // 1 allows it to be recursed once, and 2 allows it to be recursed twice and so on.
+        // A value greater than 10 is not allowed, and if a protobuf record has more depth for
+        // recursive fields than the allowed value, it will be truncated and some fields may be
+        // discarded.
+        // SQL Schema for the protobuf message `message Person { string name = 1; Person bff = 2}`
+        // will vary based on the value of "recursive.fields.max.depth".
+        // 0: struct<name: string, bff: null>
+        // 1: struct<name string, bff: <name: string, bff: null>>
+        // 2: struct<name string, bff: <name: string, bff: struct<name: string, bff: null>>> ...
+        val recordName = fd.getMessageType.getFullName
+        val recursiveDepth = existingRecordNames.getOrElse(recordName, 0)
+        val recursiveFieldMaxDepth = protobufOptions.recursiveFieldMaxDepth
+        if (existingRecordNames.contains(recordName) && (recursiveFieldMaxDepth < 0 ||
+          recursiveFieldMaxDepth > 10)) {
           throw QueryCompilationErrors.foundRecursionInProtobufSchema(fd.toString())
+        } else if (existingRecordNames.contains(recordName) &&
+          recursiveDepth > recursiveFieldMaxDepth) {
+          Some(NullType)
+        } else {
+          val newRecordNames = existingRecordNames + (recordName -> (recursiveDepth + 1))
+          Option(
+            fd.getMessageType.getFields.asScala
+              .flatMap(structFieldFor(_, newRecordNames, protobufOptions))
+              .toSeq)
+            .filter(_.nonEmpty)
+            .map(StructType.apply)
         }
-        val newRecordNames = existingRecordNames + fd.getFullName
-
-        Option(
-          fd.getMessageType.getFields.asScala
-            .flatMap(structFieldFor(_, newRecordNames))
-            .toSeq)
-          .filter(_.nonEmpty)
-          .map(StructType.apply)
       case other =>
         throw QueryCompilationErrors.protobufTypeUnsupportedYetError(other.toString)
     }

connector/protobuf/src/test/resources/protobuf/functions_suite.proto

@@ -170,4 +170,87 @@ message timeStampMsg {
 message durationMsg {
   string key = 1;
   Duration duration = 2;
-}
+}
+
+message OneOfEvent {
+  string key = 1;
+  oneof payload {
+    int32 col_1 = 2;
+    string col_2 = 3;
+    int64 col_3 = 4;
+  }
+  repeated string col_4 = 5;
+}
+
+message EventWithRecursion {
+  int32 key = 1;
+  messageA a = 2;
+}
+message messageA {
+  EventWithRecursion a = 1;
+  messageB b = 2;
+}
+message messageB {
+  EventWithRecursion aa = 1;
+  messageC c = 2;
+}
+message messageC {
+  EventWithRecursion aaa = 1;
+  int32 key= 2;
+}
+
+message Employee {
+  string firstName = 1;
+  string lastName = 2;
+  oneof role {
+    IC ic = 3;
+    EM em = 4;
+    EM2 em2 = 5;
+  }
+}
+
+message IC {
+  repeated string skills = 1;
+  Employee icManager = 2;
+}
+
+message EM {
+  int64 teamsize = 1;
+  Employee emManager = 2;
+}
+
+message EM2 {
+  int64 teamsize = 1;
+  Employee em2Manager = 2;
+}
+
+message EventPerson {
+  string name = 1;
+  EventPerson bff = 2;
+}
+
+message OneOfEventWithRecursion {
+  string key = 1;
+  oneof payload {
+    EventRecursiveA recursiveA = 3;
+    EventRecursiveB recursiveB = 6;
+  }
+  string value = 7;
+}
+
+message EventRecursiveA {
+  OneOfEventWithRecursion recursiveA = 1;
+  string key = 2;
+}
+
+message EventRecursiveB {
+  string key = 1;
+  string value = 2;
+  OneOfEventWithRecursion recursiveA = 3;
+}
+
+message Status {
+  int32 id = 1;
+  Timestamp trade_time = 2;
+  Status status = 3;
+}