Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -307,101 +307,51 @@ private[sql] class ProtobufDeserializer(
updater.set(ordinal, UTF8String.fromString(jsonStr))

// Handle well known wrapper types. We unpack the value field when the desired
// output type is a primitive (determined by the option in [[ProtobufOptions]])
// output type is a primitive (determined by the option in [[ProtobufOptions]]).
// A wrapper is unwrapped only when present: an absent singular wrapper is nulled by the
// caller before reaching here, and container elements are always present. So a present
// wrapper -- even an empty one -- carries a value, and unwrapWktValue reads the inner
// scalar (its default when unset), never null, independent of `emit.default.values`.
case (MESSAGE, BooleanType)
if protoType.getMessageType.getFullName == BoolValue.getDescriptor.getFullName =>
(updater, ordinal, value) =>
val dm = value.asInstanceOf[DynamicMessage]
val unwrapped = getFieldValue(dm, dm.getDescriptorForType.getFields.get(0))
if (unwrapped == null) {
updater.setNullAt(ordinal)
} else {
updater.setBoolean(ordinal, unwrapped.asInstanceOf[Boolean])
}
updater.setBoolean(ordinal, unwrapWktValue(value).asInstanceOf[Boolean])
case (MESSAGE, IntegerType)
if (protoType.getMessageType.getFullName == Int32Value.getDescriptor.getFullName
|| protoType.getMessageType.getFullName == UInt32Value.getDescriptor.getFullName) =>
(updater, ordinal, value) =>
val dm = value.asInstanceOf[DynamicMessage]
val unwrapped = getFieldValue(dm, dm.getDescriptorForType.getFields.get(0))
if (unwrapped == null) {
updater.setNullAt(ordinal)
} else {
updater.setInt(ordinal, unwrapped.asInstanceOf[Int])
}
updater.setInt(ordinal, unwrapWktValue(value).asInstanceOf[Int])
case (MESSAGE, LongType)
if (protoType.getMessageType.getFullName == UInt32Value.getDescriptor.getFullName) =>
(updater, ordinal, value) =>
val dm = value.asInstanceOf[DynamicMessage]
val unwrapped = getFieldValue(dm, dm.getDescriptorForType.getFields.get(0))
if (unwrapped == null) {
updater.setNullAt(ordinal)
} else {
updater.setLong(ordinal, Integer.toUnsignedLong(unwrapped.asInstanceOf[Int]))
}
updater.setLong(ordinal, Integer.toUnsignedLong(unwrapWktValue(value).asInstanceOf[Int]))
case (MESSAGE, LongType)
if (protoType.getMessageType.getFullName == Int64Value.getDescriptor.getFullName
|| protoType.getMessageType.getFullName == UInt64Value.getDescriptor.getFullName) =>
(updater, ordinal, value) =>
val dm = value.asInstanceOf[DynamicMessage]
val unwrapped = getFieldValue(dm, dm.getDescriptorForType.getFields.get(0))
if (unwrapped == null) {
updater.setNullAt(ordinal)
} else {
updater.setLong(ordinal, unwrapped.asInstanceOf[Long])
}
updater.setLong(ordinal, unwrapWktValue(value).asInstanceOf[Long])
case (MESSAGE, DecimalType.LongDecimal)
if (protoType.getMessageType.getFullName == UInt64Value.getDescriptor.getFullName) =>
(updater, ordinal, value) =>
val dm = value.asInstanceOf[DynamicMessage]
val unwrapped = getFieldValue(dm, dm.getDescriptorForType.getFields.get(0))
if (unwrapped == null) {
updater.setNullAt(ordinal)
} else {
val dec = Decimal.fromString(
UTF8String.fromString(java.lang.Long.toUnsignedString(unwrapped.asInstanceOf[Long])))
updater.setDecimal(ordinal, dec)
}
val dec = Decimal.fromString(UTF8String.fromString(
java.lang.Long.toUnsignedString(unwrapWktValue(value).asInstanceOf[Long])))
updater.setDecimal(ordinal, dec)
case (MESSAGE, StringType)
if protoType.getMessageType.getFullName == StringValue.getDescriptor.getFullName =>
(updater, ordinal, value) =>
val dm = value.asInstanceOf[DynamicMessage]
val unwrapped = getFieldValue(dm, dm.getDescriptorForType.getFields.get(0))
if (unwrapped == null) {
updater.setNullAt(ordinal)
} else {
updater.set(ordinal, UTF8String.fromString(unwrapped.asInstanceOf[String]))
}
updater.set(ordinal, UTF8String.fromString(unwrapWktValue(value).asInstanceOf[String]))
case (MESSAGE, BinaryType)
if protoType.getMessageType.getFullName == BytesValue.getDescriptor.getFullName =>
(updater, ordinal, value) =>
val dm = value.asInstanceOf[DynamicMessage]
val unwrapped = getFieldValue(dm, dm.getDescriptorForType.getFields.get(0))
if (unwrapped == null) {
updater.setNullAt(ordinal)
} else {
updater.set(ordinal, unwrapped.asInstanceOf[ByteString].toByteArray)
}
updater.set(ordinal, unwrapWktValue(value).asInstanceOf[ByteString].toByteArray)
case (MESSAGE, FloatType)
if protoType.getMessageType.getFullName == FloatValue.getDescriptor.getFullName =>
(updater, ordinal, value) =>
val dm = value.asInstanceOf[DynamicMessage]
val unwrapped = getFieldValue(dm, dm.getDescriptorForType.getFields.get(0))
if (unwrapped == null) {
updater.setNullAt(ordinal)
} else {
updater.setFloat(ordinal, unwrapped.asInstanceOf[Float])
}
updater.setFloat(ordinal, unwrapWktValue(value).asInstanceOf[Float])
case (MESSAGE, DoubleType)
if protoType.getMessageType.getFullName == DoubleValue.getDescriptor.getFullName =>
(updater, ordinal, value) =>
val dm = value.asInstanceOf[DynamicMessage]
val unwrapped = getFieldValue(dm, dm.getDescriptorForType.getFields.get(0))
if (unwrapped == null) {
updater.setNullAt(ordinal)
} else {
updater.setDouble(ordinal, unwrapped.asInstanceOf[Double])
}
updater.setDouble(ordinal, unwrapWktValue(value).asInstanceOf[Double])

case (MESSAGE, st: StructType) =>
val writeRecord = getRecordWriter(
Expand Down Expand Up @@ -520,6 +470,17 @@ private[sql] class ProtobufDeserializer(
}
}

// Reads the inner `value` scalar of a present well-known wrapper message
// (google.protobuf.{Bool,Int32,...}Value). Unlike getFieldValue, this returns the scalar's
// default (0/""/false/empty-bytes) rather than null when the inner value is unset, matching the
// proto3 wrapper semantics where a present wrapper always carries a value. Absent singular
// wrappers never reach here (they are nulled by the caller), so this is only invoked for present
// wrappers, and is independent of `emit.default.values`.
private def unwrapWktValue(value: Any): AnyRef = {
val dm = value.asInstanceOf[DynamicMessage]
dm.getField(dm.getDescriptorForType.getFields.get(0))
}

// 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))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2047,11 +2047,10 @@ class ProtobufFunctionsSuite extends SharedSparkSession with ProtobufTestBase
)
}
} else {
if (defaults == "false") {
checkAnswer(parsedExplicitZero, expectedEmpty)
} else {
checkAnswer(parsedExplicitZero, Seq((0)).toDF("int32_val"))
}
// When unwrapping, a present wrapper carries a value even if it is the inner scalar's
// default, so an explicit zero unwraps to 0 regardless of emit.default.values (which
// only governs bare proto3 scalars, not a wrapper message's presence).
checkAnswer(parsedExplicitZero, Seq((0)).toDF("int32_val"))
}

// For nonzero, we should get back the number or wrapped version regardless
Expand Down Expand Up @@ -2080,6 +2079,35 @@ class ProtobufFunctionsSuite extends SharedSparkSession with ProtobufTestBase
}
}

test("well known wrappers with empty container elements unwrap to defaults") {
// A repeated/map field of unwrapped wrappers uses a non-null container
// (containsNull = false / valueContainsNull = false). A present-but-empty wrapper element
// must unwrap to the inner scalar's default, not null -- a null in a non-null container
// crashes downstream. Other container-wrapper tests only use non-empty elements, so this
// covers the empty-element case for both a repeated and a map field.
val message = spark.range(1).select(
lit(
WellKnownWrapperTypes.newBuilder()
.addInt32List(Int32Value.getDefaultInstance) // empty element -> 0
.addInt32List(Int32Value.of(7))
.putWktMap(1, StringValue.getDefaultInstance) // empty value -> ""
.build().toByteArray
).as("raw_proto"))

val opt = Map("unwrap.primitive.wrapper.types" -> "true")
checkWithFileAndClassName("WellKnownWrapperTypes") { case (name, descFilePathOpt) =>
val parsed = message.select(
from_protobuf_wrapper($"raw_proto", name, descFilePathOpt, opt).as("proto"))

checkAnswer(
parsed.select("proto.int32_list"),
spark.range(1).select(typedLit(List(0, 7)).as("int32_list")))
checkAnswer(
parsed.select("proto.wkt_map"),
spark.range(1).select(typedLit(Map(1 -> "")).as("wkt_map")))
}
}

test("test well known wrappers with upcast ints") {
// Test that the unwrap primitives behavior and upcast uint64 work correctly together.
// We'll check the deserialization behavior under every combination of the
Expand Down