Refactor to Data Source API V2

Author: jacobfi

build.sbt

@@ -2,22 +2,19 @@ organization := "com.audienceproject"
 
 name := "spark-dynamodb"
 
-version := "0.4.4"
+version := "0.5.0"
 
 description := "Plug-and-play implementation of an Apache Spark custom data source for AWS DynamoDB."
 
-scalaVersion := "2.12.7"
-
-crossScalaVersions := Seq("2.11.12", "2.12.7")
+scalaVersion := "2.11.12"
 
 resolvers += "DynamoDBLocal" at "https://s3-us-west-2.amazonaws.com/dynamodb-local/release"
 
-libraryDependencies += "com.amazonaws" % "aws-java-sdk-sts" % "1.11.571"
-libraryDependencies += "com.amazonaws" % "aws-java-sdk-dynamodb" % "1.11.571"
+libraryDependencies += "com.amazonaws" % "aws-java-sdk-sts" % "1.11.678"
+libraryDependencies += "com.amazonaws" % "aws-java-sdk-dynamodb" % "1.11.678"
 libraryDependencies += "com.amazonaws" % "DynamoDBLocal" % "[1.11,2.0)" % "test" exclude("com.google.guava", "guava")
 
-libraryDependencies += "org.apache.spark" %% "spark-sql" % "2.4.0" % "provided"
-libraryDependencies += "com.google.guava" % "guava" % "14.0.1" % "provided"
+libraryDependencies += "org.apache.spark" %% "spark-sql" % "2.4.4" % "provided"
 
 libraryDependencies += "org.scalatest" %% "scalatest" % "3.0.5" % "test"
 

src/main/scala/com/audienceproject/spark/dynamodb/DefaultSource.scala

@@ -0,0 +1,59 @@
+package com.audienceproject.spark.dynamodb.catalyst
+
+import java.util
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.util.{ArrayData, MapData}
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.UTF8String
+
+import scala.collection.JavaConverters._
+
+object JavaConverter {
+
+    def extractRowValue(row: InternalRow, index: Int, elementType: DataType): Any = {
+        elementType match {
+            case ArrayType(innerType, _) => extractArray(row.getArray(index), innerType)
+            case MapType(keyType, valueType, _) => extractMap(row.getMap(index), keyType, valueType)
+            case StructType(fields) => mapStruct(row.getStruct(index, fields.length), fields)
+            case StringType => row.getString(index)
+            case _ => row.get(index, elementType)
+        }
+    }
+
+    def extractArray(array: ArrayData, elementType: DataType): Any = {
+        elementType match {
+            case ArrayType(innerType, _) => array.toSeq[ArrayData](elementType).map(extractArray(_, innerType)).asJava
+            case MapType(keyType, valueType, _) => array.toSeq[MapData](elementType).map(extractMap(_, keyType, valueType)).asJava
+            case structType: StructType => array.toSeq[InternalRow](structType).map(mapStruct(_, structType.fields)).asJava
+            case StringType => convertStringArray(array).asJava
+            case _ => array.toSeq[Any](elementType).asJava
+        }
+    }
+
+    def extractMap(map: MapData, keyType: DataType, valueType: DataType): util.Map[String, Any] = {
+        if (keyType != StringType) throw new IllegalArgumentException(
+            s"Invalid Map key type '${keyType.typeName}'. DynamoDB only supports String as Map key type.")
+        val keys = convertStringArray(map.keyArray())
+        val values = valueType match {
+            case ArrayType(innerType, _) => map.valueArray().toSeq[ArrayData](valueType).map(extractArray(_, innerType))
+            case MapType(innerKeyType, innerValueType, _) => map.valueArray().toSeq[MapData](valueType).map(extractMap(_, innerKeyType, innerValueType))
+            case structType: StructType => map.valueArray().toSeq[InternalRow](structType).map(mapStruct(_, structType.fields))
+            case StringType => convertStringArray(map.valueArray())
+            case _ => map.valueArray().toSeq[Any](valueType)
+        }
+        val kvPairs = for (i <- 0 until map.numElements()) yield keys(i) -> values(i)
+        Map(kvPairs: _*).asJava
+    }
+
+    def mapStruct(row: InternalRow, fields: Seq[StructField]): util.Map[String, Any] = {
+        val kvPairs = for (i <- 0 until row.numFields)
+            yield fields(i).name -> extractRowValue(row, i, fields(i).dataType)
+        Map(kvPairs: _*).asJava
+    }
+
+
+    def convertStringArray(array: ArrayData): Seq[String] =
+        array.toSeq[UTF8String](StringType).map(_.toString)
+
+}

src/main/scala/com/audienceproject/spark/dynamodb/catalyst/JavaConverter.scala

@@ -52,6 +52,7 @@ private[dynamodb] trait DynamoConnector {
                 .build()
         )
     }
+
     def getDynamoDBAsyncClient(region: Option[String] = None, roleArn: Option[String] = None): AmazonDynamoDBAsync = {
         val chosenRegion = region.getOrElse(sys.env.getOrElse("aws.dynamodb.region", "us-east-1"))
         val credentials = getCredentials(chosenRegion, roleArn)
@@ -101,7 +102,7 @@ private[dynamodb] trait DynamoConnector {
 
     val itemLimit: Int
 
-    val totalSizeInBytes: Long
+    val totalSegments: Int
 
     def scan(segmentNum: Int, columns: Seq[String], filters: Seq[Filter]): ItemCollection[ScanOutcome]
 

src/main/scala/com/audienceproject/spark/dynamodb/connector/DynamoConnector.scala

@@ -20,13 +20,13 @@
   */
 package com.audienceproject.spark.dynamodb.connector
 
-import org.apache.spark.sql.Row
+import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.types.StructType
 
 trait DynamoWritable {
 
     val writeLimit: Double
 
-    def putItems(schema: StructType, batchSize: Int)(items: Iterator[Row]): Unit
+    def putItems(schema: StructType, items: Seq[InternalRow]): Unit
 
 }