Implemented update writer. All tests now pass

Author: jacobfi

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

@@ -1,3 +1,23 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one
+  * or more contributor license agreements.  See the NOTICE file
+  * distributed with this work for additional information
+  * regarding copyright ownership.  The ASF licenses this file
+  * to you under the Apache License, Version 2.0 (the
+  * "License"); you may not use this file except in compliance
+  * with the License.  You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing,
+  * software distributed under the License is distributed on an
+  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+  * KIND, either express or implied.  See the License for the
+  * specific language governing permissions and limitations
+  * under the License.
+  *
+  * Copyright © 2019 AudienceProject. All rights reserved.
+  */
 package com.audienceproject.spark.dynamodb.catalyst
 
 import java.util
@@ -11,44 +31,44 @@ import scala.collection.JavaConverters._
 
 object JavaConverter {
 
-    def extractRowValue(row: InternalRow, index: Int, elementType: DataType): Any = {
+    def convertRowValue(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) => extractStruct(row.getStruct(index, fields.length), fields)
+            case ArrayType(innerType, _) => convertArray(row.getArray(index), innerType)
+            case MapType(keyType, valueType, _) => convertMap(row.getMap(index), keyType, valueType)
+            case StructType(fields) => convertStruct(row.getStruct(index, fields.length), fields)
             case StringType => row.getString(index)
             case _ => row.get(index, elementType)
         }
     }
 
-    def extractArray(array: ArrayData, elementType: DataType): Any = {
+    def convertArray(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(extractStruct(_, structType.fields)).asJava
+            case ArrayType(innerType, _) => array.toSeq[ArrayData](elementType).map(convertArray(_, innerType)).asJava
+            case MapType(keyType, valueType, _) => array.toSeq[MapData](elementType).map(convertMap(_, keyType, valueType)).asJava
+            case structType: StructType => array.toSeq[InternalRow](structType).map(convertStruct(_, 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] = {
+    def convertMap(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(extractStruct(_, structType.fields))
+            case ArrayType(innerType, _) => map.valueArray().toSeq[ArrayData](valueType).map(convertArray(_, innerType))
+            case MapType(innerKeyType, innerValueType, _) => map.valueArray().toSeq[MapData](valueType).map(convertMap(_, innerKeyType, innerValueType))
+            case structType: StructType => map.valueArray().toSeq[InternalRow](structType).map(convertStruct(_, 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 extractStruct(row: InternalRow, fields: Seq[StructField]): util.Map[String, Any] = {
+    def convertStruct(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)
+            yield fields(i).name -> convertRowValue(row, i, fields(i).dataType)
         Map(kvPairs: _*).asJava
     }
 

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

@@ -0,0 +1,58 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one
+  * or more contributor license agreements.  See the NOTICE file
+  * distributed with this work for additional information
+  * regarding copyright ownership.  The ASF licenses this file
+  * to you under the Apache License, Version 2.0 (the
+  * "License"); you may not use this file except in compliance
+  * with the License.  You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing,
+  * software distributed under the License is distributed on an
+  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+  * KIND, either express or implied.  See the License for the
+  * specific language governing permissions and limitations
+  * under the License.
+  *
+  * Copyright © 2019 AudienceProject. All rights reserved.
+  */
+package com.audienceproject.spark.dynamodb.connector
+
+import org.apache.spark.sql.types.{DataType, StructType}
+
+private[dynamodb] class ColumnSchema(keySchema: KeySchema,
+                                     sparkSchema: StructType) {
+
+    type Attr = (String, Int, DataType)
+
+    private val columnNames = sparkSchema.map(_.name)
+
+    private val keyIndices = keySchema match {
+        case KeySchema(hashKey, None) =>
+            val hashKeyIndex = columnNames.indexOf(hashKey)
+            val hashKeyType = sparkSchema(hashKey).dataType
+            Left(hashKey, hashKeyIndex, hashKeyType)
+        case KeySchema(hashKey, Some(rangeKey)) =>
+            val hashKeyIndex = columnNames.indexOf(hashKey)
+            val rangeKeyIndex = columnNames.indexOf(rangeKey)
+            val hashKeyType = sparkSchema(hashKey).dataType
+            val rangeKeyType = sparkSchema(rangeKey).dataType
+            Right((hashKey, hashKeyIndex, hashKeyType), (rangeKey, rangeKeyIndex, rangeKeyType))
+    }
+
+    private val attributeIndices = columnNames.zipWithIndex.filterNot({
+        case (name, _) => keySchema match {
+            case KeySchema(hashKey, None) => name == hashKey
+            case KeySchema(hashKey, Some(rangeKey)) => name == hashKey || name == rangeKey
+        }
+    }).map({
+        case (name, index) => (name, index, sparkSchema(name).dataType)
+    })
+
+    def keys(): Either[Attr, (Attr, Attr)] = keyIndices
+
+    def attributes(): Seq[Attr] = attributeIndices
+
+}

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

@@ -20,13 +20,18 @@
   */
 package com.audienceproject.spark.dynamodb.connector
 
+import com.amazonaws.services.dynamodbv2.document.DynamoDB
+import com.audienceproject.spark.dynamodb.util.RateLimiter
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.types.StructType
 
-trait DynamoWritable {
+private[dynamodb] trait DynamoWritable {
 
     val writeLimit: Double
 
-    def putItems(schema: StructType, items: Seq[InternalRow]): Unit
+    def putItems(columnSchema: ColumnSchema, items: Seq[InternalRow])
+                (client: DynamoDB, rateLimiter: RateLimiter): Unit
+
+    def updateItem(columnSchema: ColumnSchema, item: InternalRow)
+                  (client: DynamoDB, rateLimiter: RateLimiter): Unit
 
 }

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

@@ -1,3 +1,23 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one
+  * or more contributor license agreements.  See the NOTICE file
+  * distributed with this work for additional information
+  * regarding copyright ownership.  The ASF licenses this file
+  * to you under the Apache License, Version 2.0 (the
+  * "License"); you may not use this file except in compliance
+  * with the License.  You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing,
+  * software distributed under the License is distributed on an
+  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+  * KIND, either express or implied.  See the License for the
+  * specific language governing permissions and limitations
+  * under the License.
+  *
+  * Copyright © 2019 AudienceProject. All rights reserved.
+  */
 package com.audienceproject.spark.dynamodb.connector
 
 import com.amazonaws.services.dynamodbv2.model.{KeySchemaElement, KeyType}