added logic from PR 415 Azure#415

bgmarsh · Oct 28, 2020 · 74c87f8 · 74c87f8
1 parent d286a44
commit 74c87f8
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 33 deletions.
diff --git a/src/main/scala/com/microsoft/azure/cosmosdb/spark/Constants.scala b/src/main/scala/com/microsoft/azure/cosmosdb/spark/Constants.scala
@@ -23,6 +23,6 @@
 package com.microsoft.azure.cosmosdb.spark
 
 object Constants {
-  val currentVersion = "2.4.0_2.11-3.3.1"
+  val currentVersion = "3.0.1_2.12-3.5.0"
   val userAgentSuffix = s" SparkConnector/$currentVersion"
 }
diff --git a/src/main/scala/com/microsoft/azure/cosmosdb/spark/util/StreamingUtils.scala b/src/main/scala/com/microsoft/azure/cosmosdb/spark/util/StreamingUtils.scala
@@ -1,25 +1,25 @@
 /**
-  * The MIT License (MIT)
-  * Copyright (c) 2016 Microsoft Corporation
-  *
-  * Permission is hereby granted, free of charge, to any person obtaining a copy
-  * of this software and associated documentation files (the "Software"), to deal
-  * in the Software without restriction, including without limitation the rights
-  * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-  * copies of the Software, and to permit persons to whom the Software is
-  * furnished to do so, subject to the following conditions:
-  *
-  * The above copyright notice and this permission notice shall be included in all
-  * copies or substantial portions of the Software.
-  *
-  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-  * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-  * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-  * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-  * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-  * SOFTWARE.
-  */
+ * The MIT License (MIT)
+ * Copyright (c) 2016 Microsoft Corporation
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
 package org.apache.spark.sql.cosmosdb.util
 
 import java.time.temporal.ChronoUnit
@@ -31,8 +31,7 @@ import com.microsoft.azure.cosmosdb.{Document, RequestOptions, ResourceResponse}
 import com.microsoft.azure.cosmosdb.spark.config.{Config, CosmosDBConfig}
 import com.microsoft.azure.cosmosdb.spark.schema.CosmosDBRowConverter
 import com.microsoft.azure.cosmosdb.spark.streaming.CosmosDBWriteStreamRetryPolicy
-import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.encoders.RowEncoder
+import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow}
 import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.{DataFrame, SQLContext}
@@ -44,21 +43,19 @@ object StreamingUtils extends Serializable {
 
   def createDataFrameStreaming(df: DataFrame, schema: StructType, sqlContext: SQLContext): DataFrame = {
 
-    val enconder = RowEncoder.apply(schema)
-    val mappedRdd = df.rdd.map(row => {
-      enconder.toRow(row)
-    })
+    val convert = CatalystTypeConverters.createToCatalystConverter(schema)
+    val mappedRdd = df.rdd.map(convert(_).asInstanceOf[InternalRow])
     sqlContext.internalCreateDataFrame(mappedRdd, schema, isStreaming = true)
   }
 }
 
 class StreamingWriteTask extends Serializable with CosmosDBLoggingTrait {
 
   def importStreamingData[D: ClassTag](
-    iter: Iterator[D],
-    schemaOutput: Seq[Attribute],
-    config: Config,
-    retryPolicy: CosmosDBWriteStreamRetryPolicy): Unit = {
+                                        iter: Iterator[D],
+                                        schemaOutput: Seq[Attribute],
+                                        config: Config,
+                                        retryPolicy: CosmosDBWriteStreamRetryPolicy): Unit = {
 
     val upsert: Boolean = config
       .getOrElse(CosmosDBConfig.Upsert, String.valueOf(CosmosDBConfig.DefaultUpsert))
@@ -89,4 +86,4 @@ class StreamingWriteTask extends Serializable with CosmosDBLoggingTrait {
     val latency = Math.abs(ChronoUnit.MILLIS.between(LocalDateTime.now(), startTime))
     logInfo(s"Batch of $count records written with latency $latency milliseconds")
   }
-}
+}