Cosmos Spark: Fixes conversion for MapType schema (#22291)

* Missing toMap * Adding test case
Azure · Jun 16, 2021 · 77f6412 · 77f6412
1 parent 3b91c34
commit 77f6412
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/...zure-cosmos-spark_3-1_2-12/src/main/scala/com/azure/cosmos/spark/CosmosRowConverter.scala b/...zure-cosmos-spark_3-1_2-12/src/main/scala/com/azure/cosmos/spark/CosmosRowConverter.scala
@@ -324,7 +324,7 @@ private object CosmosRowConverter
             jsonNode.fields().asScala
                 .map(element => (
                     element.getKey,
-                    convertToSparkDataType(map.valueType, element.getValue, schemaConversionMode)))
+                    convertToSparkDataType(map.valueType, element.getValue, schemaConversionMode))).toMap
         case (arrayNode: ArrayNode, array: ArrayType) =>
             arrayNode.elements().asScala
               .map(convertToSparkDataType(array.elementType, _, schemaConversionMode)).toArray

diff --git a/...-cosmos-spark_3-1_2-12/src/test/scala/com/azure/cosmos/spark/CosmosRowConverterSpec.scala b/...-cosmos-spark_3-1_2-12/src/test/scala/com/azure/cosmos/spark/CosmosRowConverterSpec.scala
@@ -705,12 +705,25 @@ class CosmosRowConverterSpec extends UnitSpec with CosmosLoggingTrait {
     val nestedObjectNode: ObjectNode = objectNode.putObject(colName1)
     nestedObjectNode.put(colName1, colVal1)
     objectNode.put(colName2, colVal2)
+
+    // with struct
     val schema = StructType(Seq(StructField(colName1, StructType(Seq(StructField(colName1, StringType)))),
       StructField(colName2, StringType)))
     val row = CosmosRowConverter.fromObjectNodeToRow(schema, objectNode, SchemaConversionModes.Relaxed)
     val asStruct = row.getStruct(0)
     asStruct.getString(0) shouldEqual colVal1
     row.getString(1) shouldEqual colVal2
+
+    // with map
+    val schemaWithMap = StructType(Seq(
+      StructField(colName1, MapType(keyType = StringType, valueType = StringType, valueContainsNull = false)),
+      StructField(colName2, StringType)))
+
+    val rowWithMap = CosmosRowConverter.fromObjectNodeToRow(schemaWithMap, objectNode, SchemaConversionModes.Relaxed)
+
+    val convertedMap = rowWithMap.getMap[String, String](0)
+    convertedMap(colName1) shouldEqual colVal1
+    rowWithMap.getString(1) shouldEqual colVal2
   }
 
   "raw in ObjectNode" should "translate to Row" in {