This is an automated email from the ASF dual-hosted git repository.
gurwls223 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 25fb211ddea [SPARK-43051][CONNECT] Add option to emit default values
25fb211ddea is described below
commit 25fb211ddea90abd336300c0a3627f064297b895
Author: Parth Upadhyay <parth.upadh...@gmail.com>
AuthorDate: Thu May 4 13:08:11 2023 +0900
[SPARK-43051][CONNECT] Add option to emit default values
### JIRA
JIRA: https://issues.apache.org/jira/browse/SPARK-43051
### What changes were proposed in this pull request?
Currently, when deserializing protobufs using `from_protobuf`, fields that
are not explicitly present in the serialized message are deserialized as null
in the resulting struct. However this includes singular proto3 scalars set
explicitly to their default values, as they will [not
appear](https://protobuf.dev/programming-guides/field_presence/#presence-in-tag-value-stream-wire-format-serialization)
in the serialized protobuf.
For example, given a message format like
```
syntax = "proto3";
message Person {
string name = 1;
int64 age = 2;
optional string middle_name = 3;
optional int64 salary = 4;
}
```
and an example message like
```
SearchRequest(age = 0, middle_name = "")
```
the result from calling `from_protobuf` on the serialized form of the above
message would be
```
{"name": null, "age": null, "middle_name": "", "salary": null}
```
It can be useful to deserialize these fields as their defaults, e.g.:
```
{"name": "", "age": 0, "middle_name": "", "salary": null}
```
This behavior also exists in other major libraries, e.g.
`includingDefaultValues` in
[jsonformat](https://protobuf.dev/reference/java/api-docs/com/google/protobuf/util/JsonFormat.Printer.html#includingDefaultValueFields--),
or `emitDefaults` in
[jsonpb](https://pkg.go.dev/github.com/golang/protobuf/jsonpb#Marshaler).
In this PR I implemented this behavior by adding an option,
`emit.default.values` that can be passed to the options map for `from_protobuf`
which controls whether to materialize these defaults or not.
### Why are the changes needed?
Additional functionality to help control the deserialization behavior.
### Does this PR introduce _any_ user-facing change?
Yes, it provides a new option that can be accepted by `from_protobuf`
### How was this patch tested?
I added test cases that assert the deserialization behavior for unset and
default values for every type in proto2 and proto3.
Closes #40686 from
justaparth/parth/allow-materializing-default-protobuf-deserialization.
Authored-by: Parth Upadhyay <parth.upadh...@gmail.com>
Signed-off-by: Hyukjin Kwon <gurwls...@apache.org>
---
.../sql/protobuf/ProtobufDataToCatalyst.scala | 8 +-
.../spark/sql/protobuf/ProtobufDeserializer.scala | 34 ++-
.../spark/sql/protobuf/utils/ProtobufOptions.scala | 35 ++++
.../test/resources/protobuf/functions_suite.desc | Bin 9648 -> 10635 bytes
.../test/resources/protobuf/functions_suite.proto | 28 +++
.../test/resources/protobuf/proto2_messages.desc | Bin 288 -> 929 bytes
.../test/resources/protobuf/proto2_messages.proto | 24 +++
.../sql/protobuf/ProtobufFunctionsSuite.scala | 232 ++++++++++++++++++++-
8 files changed, 353 insertions(+), 8 deletions(-)
diff --git
a/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/ProtobufDataToCatalyst.scala
b/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/ProtobufDataToCatalyst.scala
index da44f94d5ea..cf6114f2f6c 100644
---
a/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/ProtobufDataToCatalyst.scala
+++
b/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/ProtobufDataToCatalyst.scala
@@ -21,6 +21,7 @@ import scala.util.control.NonFatal
import com.google.protobuf.DynamicMessage
+import org.apache.spark.sql.catalyst.NoopFilters
import org.apache.spark.sql.catalyst.expressions.{ExpectsInputTypes,
Expression, SpecificInternalRow, UnaryExpression}
import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext,
CodeGenerator, ExprCode}
import org.apache.spark.sql.catalyst.util.{FailFastMode, ParseMode,
PermissiveMode}
@@ -63,7 +64,12 @@ private[protobuf] case class ProtobufDataToCatalyst(
@transient private lazy val fieldsNumbers =
messageDescriptor.getFields.asScala.map(f => f.getNumber).toSet
- @transient private lazy val deserializer = new
ProtobufDeserializer(messageDescriptor, dataType)
+ @transient private lazy val deserializer =
+ new ProtobufDeserializer(
+ messageDescriptor,
+ dataType,
+ new NoopFilters,
+ protobufOptions.emitDefaultValues)
@transient private var result: DynamicMessage = _
diff --git
a/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/ProtobufDeserializer.scala
b/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/ProtobufDeserializer.scala
index 7723687a4d9..c3bc46c62ea 100644
---
a/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/ProtobufDeserializer.scala
+++
b/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/ProtobufDeserializer.scala
@@ -36,10 +36,11 @@ import org.apache.spark.unsafe.types.UTF8String
private[sql] class ProtobufDeserializer(
rootDescriptor: Descriptor,
rootCatalystType: DataType,
- filters: StructFilters) {
+ filters: StructFilters,
+ emitDefaultValues: Boolean = false) {
def this(rootDescriptor: Descriptor, rootCatalystType: DataType) = {
- this(rootDescriptor, rootCatalystType, new NoopFilters)
+ this(rootDescriptor, rootCatalystType, new NoopFilters, false)
}
private val converter: Any => Option[InternalRow] =
@@ -288,9 +289,7 @@ private[sql] class ProtobufDeserializer(
var skipRow = false
while (i < validFieldIndexes.length && !skipRow) {
val field = validFieldIndexes(i)
- val value = if (field.isRepeated || field.hasDefaultValue ||
record.hasField(field)) {
- record.getField(field)
- } else null
+ val value = getFieldValue(record, field)
fieldWriters(i)(fieldUpdater, value)
skipRow = applyFilters(i)
i += 1
@@ -299,6 +298,31 @@ private[sql] class ProtobufDeserializer(
}
}
+ private def getFieldValue(record: DynamicMessage, field: FieldDescriptor):
AnyRef = {
+ // We return a value if one of:
+ // - the field is repeated
+ // - the field is explicitly present in the serialized proto
+ // - the field is proto2 with a default
+ // - emitDefaultValues is set, and the field type is one where default
values
+ // are not present in the wire format. This includes singular proto3
scalars,
+ // but not messages / oneof / proto2.
+ // See [[ProtobufOptions]] and
https://protobuf.dev/programming-guides/field_presence
+ // for more information.
+ //
+ // Repeated fields have to be treated separately as they cannot have
`hasField`
+ // called on them.
+ if (
+ field.isRepeated
+ || record.hasField(field)
+ || field.hasDefaultValue
+ || (!field.hasPresence && this.emitDefaultValues)
+ ) {
+ record.getField(field)
+ } else {
+ null
+ }
+ }
+
// TODO: All of the code below this line is same between protobuf and avro,
it can be shared.
private def createArrayData(elementType: DataType, length: Int): ArrayData =
elementType match {
case BooleanType => UnsafeArrayData.fromPrimitiveArray(new
Array[Boolean](length))
diff --git
a/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/utils/ProtobufOptions.scala
b/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/utils/ProtobufOptions.scala
index 53036668ebf..0af8c50dc5a 100644
---
a/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/utils/ProtobufOptions.scala
+++
b/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/utils/ProtobufOptions.scala
@@ -46,6 +46,41 @@ private[sql] class ProtobufOptions(
// record has more depth than the allowed value for recursive fields, it
will be truncated
// and corresponding fields are ignored (dropped).
val recursiveFieldMaxDepth: Int =
parameters.getOrElse("recursive.fields.max.depth", "-1").toInt
+
+ // Whether to render fields with zero values when deserializing Protobufs to
a Spark struct.
+ // When a field is empty in the serialized Protobuf, this library will
deserialize them as
+ // null by default. However, this flag can control whether to render the
type-specific zero value.
+ // This operates similarly to `includingDefaultValues` in
protobuf-java-util's JsonFormat, or
+ // `emitDefaults` in golang/protobuf's jsonpb.
+ //
+ // As an example:
+ // ```
+ // syntax = "proto3";
+ // message Person {
+ // string name = 1;
+ // int64 age = 2;
+ // optional string middle_name = 3;
+ // optional int64 salary = 4;
+ // }
+ // ```
+ //
+ // And we have a proto constructed like:
+ // `Person(age=0, middle_name="")
+ //
+ // The result after calling from_protobuf without this flag set would be:
+ // `{"name": null, "age": null, "middle_name": "", "salary": null}`
+ // (age is null because zero-value singular fields are not in the wire
format in proto3).
+ //
+ //
+ // With this flag it would be:
+ // `{"name": "", "age": 0, "middle_name": "", "salary": null}`
+ //
+ // Ref: https://protobuf.dev/programming-guides/proto3/#default for
information about
+ // type-specific defaults.
+ // Ref: https://protobuf.dev/programming-guides/field_presence/ for
information about
+ // what information is available in a serialized proto.
+ val emitDefaultValues: Boolean =
+ parameters.getOrElse("emit.default.values", false.toString).toBoolean
}
private[sql] object ProtobufOptions {
diff --git
a/connector/protobuf/src/test/resources/protobuf/functions_suite.desc
b/connector/protobuf/src/test/resources/protobuf/functions_suite.desc
index 467b9cac969..80b78cd75f6 100644
Binary files
a/connector/protobuf/src/test/resources/protobuf/functions_suite.desc and
b/connector/protobuf/src/test/resources/protobuf/functions_suite.desc differ
diff --git
a/connector/protobuf/src/test/resources/protobuf/functions_suite.proto
b/connector/protobuf/src/test/resources/protobuf/functions_suite.proto
index d83ba6a4f6e..2e9add4987c 100644
--- a/connector/protobuf/src/test/resources/protobuf/functions_suite.proto
+++ b/connector/protobuf/src/test/resources/protobuf/functions_suite.proto
@@ -283,3 +283,31 @@ message Status {
Timestamp trade_time = 2;
Status status = 3;
}
+
+// Contains a representative sample of all types, using the groupings defined
+// here:
https://protobuf.dev/programming-guides/field_presence/#presence-in-proto3-apis
+message Proto3AllTypes {
+ enum NestedEnum {
+ NOTHING = 0;
+ FIRST = 1;
+ SECOND = 2;
+ }
+
+ int64 int = 1;
+ string text = 2;
+ NestedEnum enum_val = 3;
+ OtherExample message = 4;
+
+ optional int64 optional_int = 5;
+ optional string optional_text = 6;
+ optional NestedEnum optional_enum_val = 7;
+ optional OtherExample optional_message = 8;
+
+ repeated int64 repeated_num = 9;
+ repeated OtherExample repeated_message = 10;
+ oneof payload {
+ int32 option_a = 11;
+ string option_b = 12;
+ }
+ map<string, string> map = 13;
+}
diff --git
a/connector/protobuf/src/test/resources/protobuf/proto2_messages.desc
b/connector/protobuf/src/test/resources/protobuf/proto2_messages.desc
index a9e4099a7f2..eaf93670883 100644
Binary files
a/connector/protobuf/src/test/resources/protobuf/proto2_messages.desc and
b/connector/protobuf/src/test/resources/protobuf/proto2_messages.desc differ
diff --git
a/connector/protobuf/src/test/resources/protobuf/proto2_messages.proto
b/connector/protobuf/src/test/resources/protobuf/proto2_messages.proto
index a5d09df8514..75525715a3d 100644
--- a/connector/protobuf/src/test/resources/protobuf/proto2_messages.proto
+++ b/connector/protobuf/src/test/resources/protobuf/proto2_messages.proto
@@ -31,3 +31,27 @@ message FoobarWithRequiredFieldBar {
message NestedFoobarWithRequiredFieldBar {
optional FoobarWithRequiredFieldBar nested_foobar = 1;
}
+
+// Contains a representative sample of all types, using the groupings defined
+// here:
https://protobuf.dev/programming-guides/field_presence/#presence-in-proto3-apis
+message Proto2AllTypes {
+ enum NestedEnum {
+ NOTHING = 0;
+ FIRST = 1;
+ SECOND = 2;
+ }
+
+ optional int64 int = 1;
+ optional string text = 2;
+ optional NestedEnum enum_val = 3;
+ optional FoobarWithRequiredFieldBar message = 4;
+
+ repeated int64 repeated_num = 5;
+ repeated FoobarWithRequiredFieldBar repeated_message = 6;
+
+ oneof payload {
+ int32 option_a = 11;
+ string option_b = 12;
+ }
+ map<string, string> map = 13;
+}
diff --git
a/connector/protobuf/src/test/scala/org/apache/spark/sql/protobuf/ProtobufFunctionsSuite.scala
b/connector/protobuf/src/test/scala/org/apache/spark/sql/protobuf/ProtobufFunctionsSuite.scala
index 92c3c27bfae..3105d9dc8b5 100644
---
a/connector/protobuf/src/test/scala/org/apache/spark/sql/protobuf/ProtobufFunctionsSuite.scala
+++
b/connector/protobuf/src/test/scala/org/apache/spark/sql/protobuf/ProtobufFunctionsSuite.scala
@@ -19,12 +19,13 @@ package org.apache.spark.sql.protobuf
import java.sql.Timestamp
import java.time.Duration
- import scala.collection.JavaConverters._
+import scala.collection.JavaConverters._
import com.google.protobuf.{ByteString, DynamicMessage}
import org.apache.spark.sql.{AnalysisException, Column, DataFrame, QueryTest,
Row}
-import org.apache.spark.sql.functions.{lit, struct}
+import org.apache.spark.sql.functions.{lit, struct, typedLit}
+import org.apache.spark.sql.protobuf.protos.Proto2Messages.Proto2AllTypes
import org.apache.spark.sql.protobuf.protos.SimpleMessageProtos._
import
org.apache.spark.sql.protobuf.protos.SimpleMessageProtos.SimpleMessageRepeated.NestedEnum
import org.apache.spark.sql.protobuf.utils.ProtobufUtils
@@ -39,6 +40,9 @@ class ProtobufFunctionsSuite extends QueryTest with
SharedSparkSession with Prot
val testFileDesc = testFile("functions_suite.desc",
"protobuf/functions_suite.desc")
private val javaClassNamePrefix =
"org.apache.spark.sql.protobuf.protos.SimpleMessageProtos$"
+ val proto2FileDesc = testFile("proto2_messages.desc",
"protobuf/proto_messages.desc")
+ private val proto2JavaClassNamePrefix =
"org.apache.spark.sql.protobuf.protos.Proto2Messages$"
+
private def emptyBinaryDF = Seq(Array[Byte]()).toDF("binary")
/**
@@ -54,6 +58,16 @@ class ProtobufFunctionsSuite extends QueryTest with
SharedSparkSession with Prot
}
}
+ private def checkWithProto2FileAndClassName(messageName: String)(
+ fn: (String, Option[String]) => Unit): Unit = {
+ withClue("(With descriptor file)") {
+ fn(messageName, Some(proto2FileDesc))
+ }
+ withClue("(With Java class name)") {
+ fn(s"$proto2JavaClassNamePrefix$messageName", None)
+ }
+ }
+
// A wrapper to invoke the right variable of from_protobuf() depending on
arguments.
private def from_protobuf_wrapper(
col: Column,
@@ -1106,6 +1120,220 @@ class ProtobufFunctionsSuite extends QueryTest with
SharedSparkSession with Prot
}
}
+ test("test explicitly set zero values - proto3") {
+ // All fields explicitly zero. Message, map, repeated, and oneof fields
+ // are left unset, as null is their zero value.
+ val explicitZero = spark.range(1).select(
+ lit(
+ Proto3AllTypes.newBuilder()
+ .setInt(0)
+ .setText("")
+ .setEnumVal(Proto3AllTypes.NestedEnum.NOTHING)
+ .setOptionalInt(0)
+ .setOptionalText("")
+ .setOptionalEnumVal(Proto3AllTypes.NestedEnum.NOTHING)
+ .setOptionA(0)
+ .build()
+ .toByteArray).as("raw_proto"))
+
+ // By default, we deserialize zero values for fields without
+ // field presence (i.e. most primitives in proto3) as null.
+ // For fields with field presence, (explicitly optional, oneof, etc)
+ // we're able to get the explicitly set zero value.
+ val expected = spark.range(1).select(
+ struct(
+ lit(null).as("int"),
+ lit(null).as("text"),
+ lit(null).as("enum_val"),
+ lit(null).as("message"),
+ lit(0).as("optional_int"),
+ lit("").as("optional_text"),
+ lit("NOTHING").as("optional_enum_val"),
+ lit(null).as("optional_message"),
+ lit(Array.emptyIntArray).as("repeated_num"),
+ lit(Array.emptyByteArray).as("repeated_message"),
+ lit(0).as("option_a"),
+ lit(null).as("option_b"),
+ typedLit(Map.empty[String, String]).as("map")
+ ).as("proto")
+ )
+
+ // With the emit.default.values flag set, we'll fill in
+ // the fields without presence info.
+ val expectedWithFlag = spark.range(1).select(
+ struct(
+ lit(0).as("int"),
+ lit("").as("text"),
+ lit("NOTHING").as("enum_val"),
+ lit(null).as("message"),
+ lit(0).as("optional_int"),
+ lit("").as("optional_text"),
+ lit("NOTHING").as("optional_enum_val"),
+ lit(null).as("optional_message"),
+ lit(Array.emptyIntArray).as("repeated_num"),
+ lit(Array.emptyByteArray).as("repeated_message"),
+ lit(0).as("option_a"),
+ lit(null).as("option_b"),
+ typedLit(Map.empty[String, String]).as("map")
+ ).as("proto")
+ )
+
+ checkWithFileAndClassName("Proto3AllTypes") { case (name, descFilePathOpt)
=>
+ checkAnswer(
+ explicitZero.select(
+ from_protobuf_wrapper($"raw_proto", name,
descFilePathOpt).as("proto")),
+ expected)
+ checkAnswer(
+ explicitZero.select(from_protobuf_wrapper(
+ $"raw_proto",
+ name,
+ descFilePathOpt,
+ Map("emit.default.values" -> "true")).as("proto")),
+ expectedWithFlag)
+ }
+ }
+
+ test("test unset values - proto3") {
+ // Test how we deserialize fields not being present at all.
+ val empty = spark.range(1)
+ .select(lit(
+ Proto3AllTypes.newBuilder().build().toByteArray
+ ).as("raw_proto"))
+
+ val expected = spark.range(1).select(
+ struct(
+ lit(null).as("int"),
+ lit(null).as("text"),
+ lit(null).as("enum_val"),
+ lit(null).as("message"),
+ lit(null).as("optional_int"),
+ lit(null).as("optional_text"),
+ lit(null).as("optional_enum_val"),
+ lit(null).as("optional_message"),
+ lit(Array.emptyIntArray).as("repeated_num"),
+ lit(Array.emptyByteArray).as("repeated_message"),
+ lit(null).as("option_a"),
+ lit(null).as("option_b"),
+ typedLit(Map.empty[String, String]).as("map")
+ ).as("proto")
+ )
+
+ // With the emit.default.values flag set, we'll fill in
+ // the fields without presence info.
+ val expectedWithFlag = spark.range(1).select(
+ struct(
+ lit(0).as("int"),
+ lit("").as("text"),
+ lit("NOTHING").as("enum_val"),
+ lit(null).as("message"),
+ lit(null).as("optional_int"),
+ lit(null).as("optional_text"),
+ lit(null).as("optional_enum_val"),
+ lit(null).as("optional_message"),
+ lit(Array.emptyIntArray).as("repeated_num"),
+ lit(Array.emptyByteArray).as("repeated_message"),
+ lit(null).as("option_a"),
+ lit(null).as("option_b"),
+ typedLit(Map.empty[String, String]).as("map")
+ ).as("proto")
+ )
+
+ checkWithFileAndClassName("Proto3AllTypes") { case (name, descFilePathOpt)
=>
+ checkAnswer(
+ empty.select(
+ from_protobuf_wrapper($"raw_proto", name,
descFilePathOpt).as("proto")),
+ expected)
+ checkAnswer(
+ empty.select(from_protobuf_wrapper(
+ $"raw_proto",
+ name,
+ descFilePathOpt,
+ Map("emit.default.values" -> "true")).as("proto")),
+ expectedWithFlag)
+ }
+ }
+
+ test("test explicitly set zero values - proto2") {
+ // All fields explicitly zero. Message, map, repeated, and oneof fields
+ // are left unset, as null is their zero value.
+ val explicitZero = spark.range(1).select(
+ lit(
+ Proto2AllTypes.newBuilder()
+ .setInt(0)
+ .setText("")
+ .setEnumVal(Proto2AllTypes.NestedEnum.NOTHING)
+ .setOptionA(0)
+ .build()
+ .toByteArray).as("raw_proto"))
+
+ // We are able to get the zero value back when deserializing since
+ // most proto2 fields have field presence information.
+ val expected = spark.range(1).select(
+ struct(
+ lit(0).as("int"),
+ lit("").as("text"),
+ lit("NOTHING").as("enum_val"),
+ lit(null).as("message"),
+ lit(Array.emptyIntArray).as("repeated_num"),
+ lit(Array.emptyByteArray).as("repeated_message"),
+ lit(0).as("option_a"),
+ lit(null).as("option_b"),
+ typedLit(Map.empty[String, String]).as("map")
+ ).as("proto")
+ )
+
+ checkWithProto2FileAndClassName("Proto2AllTypes") { case (name,
descFilePathOpt) =>
+ checkAnswer(
+ explicitZero.select(
+ from_protobuf_wrapper($"raw_proto", name,
descFilePathOpt).as("proto")),
+ expected)
+ checkAnswer(
+ explicitZero.select(from_protobuf_wrapper(
+ $"raw_proto",
+ name,
+ descFilePathOpt,
+ Map("emit.default.values" -> "true")).as("proto")),
+ expected)
+ }
+ }
+
+ test("test unset fields - proto2 types") {
+ // All fields explicitly zero. Message, map, repeated, and oneof fields
+ // have null, i.e. not set, as their empty versions.
+ val empty = spark.range(1).select(
+ lit(Proto2AllTypes.newBuilder().build().toByteArray).as("raw_proto"))
+
+ val expected = spark.range(1).select(
+ struct(
+ lit(null).as("int"),
+ lit(null).as("text"),
+ lit(null).as("enum_val"),
+ lit(null).as("message"),
+ lit(Array.emptyIntArray).as("repeated_num"),
+ lit(Array.emptyByteArray).as("repeated_message"),
+ lit(null).as("option_a"),
+ lit(null).as("option_b"),
+ typedLit(Map.empty[String, String]).as("map")
+ ).as("proto")
+ )
+
+ // emit.default.values will not materialize values for fields with presence
+ // info available, which is most fields within proto2.
+ checkWithProto2FileAndClassName("Proto2AllTypes") { case (name,
descFilePathOpt) =>
+ checkAnswer(
+ empty.select(
+ from_protobuf_wrapper($"raw_proto", name,
descFilePathOpt).as("proto")),
+ expected)
+ checkAnswer(
+ empty.select(from_protobuf_wrapper(
+ $"raw_proto",
+ name,
+ descFilePathOpt,
+ Map("emit.default.values" -> "true")).as("proto")),
+ expected)
+ }
+ }
+
def testFromProtobufWithOptions(
df: DataFrame,
expectedDf: DataFrame,
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@spark.apache.org
For additional commands, e-mail: commits-h...@spark.apache.org
