Fixed empty batch write. Fixed NPE on nulls in structs. Fixed acquire zero capacity on rate limiter.

Author: jacobfi

build.sbt

@@ -2,7 +2,7 @@ organization := "com.audienceproject"
 
 name := "spark-dynamodb"
 
-version := "1.0.0"
+version := "1.0.1"
 
 description := "Plug-and-play implementation of an Apache Spark custom data source for AWS DynamoDB."
 

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

@@ -67,8 +67,9 @@ object JavaConverter {
     }
 
     def convertStruct(row: InternalRow, fields: Seq[StructField]): util.Map[String, Any] = {
-        val kvPairs = for (i <- 0 until row.numFields)
-            yield fields(i).name -> convertRowValue(row, i, fields(i).dataType)
+        val kvPairs = for (i <- 0 until row.numFields) yield
+            if (row.isNullAt(i)) fields(i).name -> null
+            else fields(i).name -> convertRowValue(row, i, fields(i).dataType)
         Map(kvPairs: _*).asJava
     }
 

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

@@ -168,7 +168,7 @@ private[dynamodb] class TableConnector(tableName: String, parallelism: Int, para
         // Update item and rate limit on write capacity.
         val response = client.getTable(tableName).updateItem(updateItemSpec)
         Option(response.getUpdateItemResult.getConsumedCapacity)
-            .foreach(cap => rateLimiter.acquire(cap.getCapacityUnits.toInt))
+            .foreach(cap => rateLimiter.acquire(cap.getCapacityUnits.toInt max 1))
     }
 
     @tailrec
@@ -178,7 +178,7 @@ private[dynamodb] class TableConnector(tableName: String, parallelism: Int, para
         if (response.getBatchWriteItemResult.getConsumedCapacity != null) {
             response.getBatchWriteItemResult.getConsumedCapacity.asScala.map(cap => {
                 cap.getTableName -> cap.getCapacityUnits.toInt
-            }).toMap.get(tableName).foreach(units => rateLimiter.acquire(units))
+            }).toMap.get(tableName).foreach(units => rateLimiter.acquire(units max 1))
         }
         // Retry unprocessed items.
         if (response.getUnprocessedItems != null && !response.getUnprocessedItems.isEmpty) {

src/main/scala/com/audienceproject/spark/dynamodb/datasource/DynamoBatchWriter.scala

@@ -52,8 +52,10 @@ class DynamoBatchWriter(batchSize: Int,
     override def abort(): Unit = {}
 
     private def flush(): Unit = {
-        connector.putItems(columnSchema, buffer)(client, rateLimiter)
-        buffer.clear()
+        if (buffer.nonEmpty) {
+            connector.putItems(columnSchema, buffer)(client, rateLimiter)
+            buffer.clear()
+        }
     }
 
 }

src/main/scala/com/audienceproject/spark/dynamodb/datasource/ScanPartition.scala

@@ -90,7 +90,7 @@ class ScanPartition(schema: StructType,
         private def nextPage(): Unit = {
             val page = pageIterator.next()
             val result = page.getLowLevelResult
-            Option(result.getScanResult.getConsumedCapacity).foreach(cap => rateLimiter.acquire(cap.getCapacityUnits.toInt))
+            Option(result.getScanResult.getConsumedCapacity).foreach(cap => rateLimiter.acquire(cap.getCapacityUnits.toInt max 1))
             innerIterator = result.getItems.iterator().asScala.map(itemToRow(requiredColumns))
         }
 

src/test/scala/com/audienceproject/spark/dynamodb/NullValuesTest.scala

@@ -0,0 +1,44 @@
+package com.audienceproject.spark.dynamodb
+
+import com.amazonaws.services.dynamodbv2.model.{AttributeDefinition, CreateTableRequest, KeySchemaElement, ProvisionedThroughput}
+import com.audienceproject.spark.dynamodb.implicits._
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
+
+class NullValuesTest extends AbstractInMemoryTest {
+
+    test("Insert nested StructType with null values") {
+        dynamoDB.createTable(new CreateTableRequest()
+            .withTableName("NullTest")
+            .withAttributeDefinitions(new AttributeDefinition("name", "S"))
+            .withKeySchema(new KeySchemaElement("name", "HASH"))
+            .withProvisionedThroughput(new ProvisionedThroughput(5L, 5L)))
+
+        val schema = StructType(
+            Seq(
+                StructField("name", StringType, nullable = false),
+                StructField("info", StructType(
+                    Seq(
+                        StructField("age", IntegerType, nullable = true),
+                        StructField("address", StringType, nullable = true)
+                    )
+                ), nullable = true)
+            )
+        )
+
+        val rows = spark.sparkContext.parallelize(Seq(
+            Row("one", Row(30, "Somewhere")),
+            Row("two", null),
+            Row("three", Row(null, null))
+        ))
+
+        val newItemsDs = spark.createDataFrame(rows, schema)
+
+        newItemsDs.write.dynamodb("NullTest")
+
+        val validationDs = spark.read.dynamodb("NullTest")
+
+        validationDs.show(false)
+    }
+
+}