update the code to correctly handle SAM types

Author: maasg

src/main/scala/com/microsoft/azure/cosmosdb/spark/rdd/CosmosDBRDD.scala

@@ -31,6 +31,7 @@ import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.sources.Filter
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.{DataFrame, Dataset, SparkSession}
+import org.apache.spark.util.TaskCompletionListener
 import org.apache.spark.{Partition, TaskContext}
 
 import scala.collection.mutable
@@ -117,9 +118,10 @@ class CosmosDBRDD(
         var cosmosDBPartition: CosmosDBPartition = partition.asInstanceOf[CosmosDBPartition]
         logInfo(s"CosmosDBRDD:compute: Start CosmosDBRDD compute task for partition key range id ${cosmosDBPartition.partitionKeyRangeId}")
 
-        context.addTaskCompletionListener((ctx: TaskContext) => {
+        val taskCompletionListener:TaskCompletionListener = (ctx: TaskContext) => {
           logInfo(s"CosmosDBRDD:compute: CosmosDBRDD compute task completed for partition key range id ${cosmosDBPartition.partitionKeyRangeId}")
-        })
+        }
+        context.addTaskCompletionListener(taskCompletionListener)
 
         new CosmosDBRDDIterator(
           hadoopConfig,

src/main/scala/com/microsoft/azure/cosmosdb/spark/rdd/CosmosDBRDDIterator.scala

@@ -27,7 +27,7 @@ import java.util.concurrent.ConcurrentHashMap
 
 import com.fasterxml.jackson.databind.ObjectMapper
 import com.microsoft.azure.cosmosdb.internal.HttpConstants.StatusCodes
-import com.microsoft.azure.cosmosdb.internal.directconnectivity.{ ServiceUnavailableException}
+import com.microsoft.azure.cosmosdb.internal.directconnectivity.ServiceUnavailableException
 import com.microsoft.azure.cosmosdb.spark.config.{Config, CosmosDBConfig}
 import com.microsoft.azure.cosmosdb.spark.partitioner.CosmosDBPartition
 import com.microsoft.azure.cosmosdb.spark.schema._
@@ -38,6 +38,7 @@ import com.microsoft.azure.documentdb.internal.HttpConstants.SubStatusCodes
 import org.apache.commons.lang3.StringUtils
 import org.apache.spark._
 import org.apache.spark.sql.sources.Filter
+import org.apache.spark.util.TaskCompletionListener
 
 import scala.collection.mutable
 
@@ -414,10 +415,11 @@ class CosmosDBRDDIterator(hadoopConfig: mutable.Map[String, String],
     })
 
     // Register an on-task-completion callback to close the input stream.
-    taskContext.addTaskCompletionListener((context: TaskContext) => {
+    val taskCompletionListener: TaskCompletionListener = (context: TaskContext) => {
       connection.reinitializeClient()
       closeIfNeeded()
-    })
+    }
+    taskContext.addTaskCompletionListener(taskCompletionListener)
 
     if (!readingChangeFeed) {
       queryDocuments