Skip to content

Commit ce10545

Browse files
ueshincloud-fan
authored andcommitted
[SPARK-21300][SQL] ExternalMapToCatalyst should null-check map key prior to converting to internal value.
## What changes were proposed in this pull request? `ExternalMapToCatalyst` should null-check map key prior to converting to internal value to throw an appropriate Exception instead of something like NPE. ## How was this patch tested? Added a test and existing tests. Author: Takuya UESHIN <[email protected]> Closes #18524 from ueshin/issues/SPARK-21300.
1 parent de14086 commit ce10545

File tree

4 files changed

+24
-2
lines changed

4 files changed

+24
-2
lines changed

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/JavaTypeInference.scala

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -418,6 +418,7 @@ object JavaTypeInference {
418418
inputObject,
419419
ObjectType(keyType.getRawType),
420420
serializerFor(_, keyType),
421+
keyNullable = true,
421422
ObjectType(valueType.getRawType),
422423
serializerFor(_, valueType),
423424
valueNullable = true

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -494,6 +494,7 @@ object ScalaReflection extends ScalaReflection {
494494
inputObject,
495495
dataTypeFor(keyType),
496496
serializerFor(_, keyType, keyPath, seenTypeSet),
497+
keyNullable = !keyType.typeSymbol.asClass.isPrimitive,
497498
dataTypeFor(valueType),
498499
serializerFor(_, valueType, valuePath, seenTypeSet),
499500
valueNullable = !valueType.typeSymbol.asClass.isPrimitive)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects/objects.scala

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -841,18 +841,21 @@ object ExternalMapToCatalyst {
841841
inputMap: Expression,
842842
keyType: DataType,
843843
keyConverter: Expression => Expression,
844+
keyNullable: Boolean,
844845
valueType: DataType,
845846
valueConverter: Expression => Expression,
846847
valueNullable: Boolean): ExternalMapToCatalyst = {
847848
val id = curId.getAndIncrement()
848849
val keyName = "ExternalMapToCatalyst_key" + id
850+
val keyIsNull = "ExternalMapToCatalyst_key_isNull" + id
849851
val valueName = "ExternalMapToCatalyst_value" + id
850852
val valueIsNull = "ExternalMapToCatalyst_value_isNull" + id
851853

852854
ExternalMapToCatalyst(
853855
keyName,
856+
keyIsNull,
854857
keyType,
855-
keyConverter(LambdaVariable(keyName, "false", keyType, false)),
858+
keyConverter(LambdaVariable(keyName, keyIsNull, keyType, keyNullable)),
856859
valueName,
857860
valueIsNull,
858861
valueType,
@@ -868,6 +871,8 @@ object ExternalMapToCatalyst {
868871
*
869872
* @param key the name of the map key variable that used when iterate the map, and used as input for
870873
* the `keyConverter`
874+
* @param keyIsNull the nullability of the map key variable that used when iterate the map, and
875+
* used as input for the `keyConverter`
871876
* @param keyType the data type of the map key variable that used when iterate the map, and used as
872877
* input for the `keyConverter`
873878
* @param keyConverter A function that take the `key` as input, and converts it to catalyst format.
@@ -883,6 +888,7 @@ object ExternalMapToCatalyst {
883888
*/
884889
case class ExternalMapToCatalyst private(
885890
key: String,
891+
keyIsNull: String,
886892
keyType: DataType,
887893
keyConverter: Expression,
888894
value: String,
@@ -913,6 +919,7 @@ case class ExternalMapToCatalyst private(
913919

914920
val keyElementJavaType = ctx.javaType(keyType)
915921
val valueElementJavaType = ctx.javaType(valueType)
922+
ctx.addMutableState("boolean", keyIsNull, "")
916923
ctx.addMutableState(keyElementJavaType, key, "")
917924
ctx.addMutableState("boolean", valueIsNull, "")
918925
ctx.addMutableState(valueElementJavaType, value, "")
@@ -950,6 +957,12 @@ case class ExternalMapToCatalyst private(
950957
defineEntries -> defineKeyValue
951958
}
952959

960+
val keyNullCheck = if (ctx.isPrimitiveType(keyType)) {
961+
s"$keyIsNull = false;"
962+
} else {
963+
s"$keyIsNull = $key == null;"
964+
}
965+
953966
val valueNullCheck = if (ctx.isPrimitiveType(valueType)) {
954967
s"$valueIsNull = false;"
955968
} else {
@@ -972,6 +985,7 @@ case class ExternalMapToCatalyst private(
972985
$defineEntries
973986
while($entries.hasNext()) {
974987
$defineKeyValue
988+
$keyNullCheck
975989
$valueNullCheck
976990

977991
${genKeyConverter.code}

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -355,12 +355,18 @@ class ExpressionEncoderSuite extends PlanTest with AnalysisTest {
355355
checkNullable[String](true)
356356
}
357357

358-
test("null check for map key") {
358+
test("null check for map key: String") {
359359
val encoder = ExpressionEncoder[Map[String, Int]]()
360360
val e = intercept[RuntimeException](encoder.toRow(Map(("a", 1), (null, 2))))
361361
assert(e.getMessage.contains("Cannot use null as map key"))
362362
}
363363

364+
test("null check for map key: Integer") {
365+
val encoder = ExpressionEncoder[Map[Integer, String]]()
366+
val e = intercept[RuntimeException](encoder.toRow(Map((1, "a"), (null, "b"))))
367+
assert(e.getMessage.contains("Cannot use null as map key"))
368+
}
369+
364370
private def encodeDecodeTest[T : ExpressionEncoder](
365371
input: T,
366372
testName: String): Unit = {

0 commit comments

Comments
 (0)