Added DynamoBatchDeleter

Separated delete logic from putItems Restored test debugging help

Author: rhelmstetter

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

@@ -29,9 +29,12 @@ private[dynamodb] trait DynamoWritable {
  val writeLimit: Double
 
  def putItems(columnSchema: ColumnSchema, items: Seq[InternalRow])
- (client: DynamoDB, rateLimiter: RateLimiter, delete: Boolean): Unit
+ (client: DynamoDB, rateLimiter: RateLimiter): Unit
 
  def updateItem(columnSchema: ColumnSchema, item: InternalRow)
  (client: DynamoDB, rateLimiter: RateLimiter): Unit
 
+ def deleteItems(columnSchema: ColumnSchema, itema: Seq[InternalRow])
+ (client: DynamoDB, rateLimiter: RateLimiter): Unit
+
 }

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

@@ -110,56 +110,34 @@ private[dynamodb] class TableConnector(tableName: String, parallelism: Int, para
  }
 
  override def putItems(columnSchema: ColumnSchema, items: Seq[InternalRow])
- (client: DynamoDB, rateLimiter: RateLimiter, delete: Boolean): Unit = {
+ (client: DynamoDB, rateLimiter: RateLimiter): Unit = {
  // For each batch.
  val batchWriteItemSpec = new BatchWriteItemSpec().withReturnConsumedCapacity(ReturnConsumedCapacity.TOTAL)
-
- val tableWriteItems = new TableWriteItems(tableName)
- val tableWriteItemsWithItems: TableWriteItems = if (delete) {
- // check if hash key only or also range key
- columnSchema.keys() match {
- case Left((hashKey, hashKeyIndex, hashKeyType)) =>
- val hashKeys = items.map(row =>
- JavaConverter.convertRowValue(row, hashKeyIndex, hashKeyType).asInstanceOf[AnyRef])
- tableWriteItems.withHashOnlyKeysToDelete(hashKey, hashKeys: _*)
-
- case Right(((hashKey, hashKeyIndex, hashKeyType), (rangeKey, rangeKeyIndex, rangeKeyType))) =>
- val alternatingHashAndRangeKeys = items.flatMap { case row =>
+ batchWriteItemSpec.withTableWriteItems(new TableWriteItems(tableName).withItemsToPut(
+ // Map the items.
+ items.map(row => {
+ val item = new Item()
+
+ // Map primary key.
+ columnSchema.keys() match {
+ case Left((hashKey, hashKeyIndex, hashKeyType)) =>
+ item.withPrimaryKey(hashKey, JavaConverter.convertRowValue(row, hashKeyIndex, hashKeyType))
+ case Right(((hashKey, hashKeyIndex, hashKeyType), (rangeKey, rangeKeyIndex, rangeKeyType))) =>
  val hashKeyValue = JavaConverter.convertRowValue(row, hashKeyIndex, hashKeyType)
  val rangeKeyValue = JavaConverter.convertRowValue(row, rangeKeyIndex, rangeKeyType)
- Seq(hashKeyValue.asInstanceOf[AnyRef], rangeKeyValue.asInstanceOf[AnyRef])
- }
- tableWriteItems.withHashAndRangeKeysToDelete(hashKey, rangeKey, alternatingHashAndRangeKeys: _*)
- }
- } else {
- // Map the items.
- tableWriteItems.withItemsToPut(
- items.map(row => {
- val item = new Item()
-
- // Map primary key.
- columnSchema.keys() match {
- case Left((hashKey, hashKeyIndex, hashKeyType)) =>
- item.withPrimaryKey(hashKey, JavaConverter.convertRowValue(row, hashKeyIndex, hashKeyType))
- case Right(((hashKey, hashKeyIndex, hashKeyType), (rangeKey, rangeKeyIndex, rangeKeyType))) =>
- val hashKeyValue = JavaConverter.convertRowValue(row, hashKeyIndex, hashKeyType)
- val rangeKeyValue = JavaConverter.convertRowValue(row, rangeKeyIndex, rangeKeyType)
- item.withPrimaryKey(hashKey, hashKeyValue, rangeKey, rangeKeyValue)
- }
+ item.withPrimaryKey(hashKey, hashKeyValue, rangeKey, rangeKeyValue)
+ }
 
-  // Map remaining columns.
-  columnSchema.attributes().foreach({
-  case (name, index, dataType) if !row.isNullAt(index) =>
-  item.`with`(name, JavaConverter.convertRowValue(row, index, dataType))
-  case _ =>
-  })
+ // Map remaining columns.
+ columnSchema.attributes().foreach({
+ case (name, index, dataType) if !row.isNullAt(index) =>
+ item.`with`(name, JavaConverter.convertRowValue(row, index, dataType))
+ case _ =>
+ })
 
- item
- }): _*
- )
- }
-
- batchWriteItemSpec.withTableWriteItems(tableWriteItemsWithItems)
+ item
+ }): _*
+ ))
 
  val response = client.batchWriteItem(batchWriteItemSpec)
  handleBatchWriteResponse(client, rateLimiter)(response)
@@ -193,6 +171,35 @@ private[dynamodb] class TableConnector(tableName: String, parallelism: Int, para
  .foreach(cap => rateLimiter.acquire(cap.getCapacityUnits.toInt max 1))
  }
 
+ override def deleteItems(columnSchema: ColumnSchema, items: Seq[InternalRow])
+ (client: DynamoDB, rateLimiter: RateLimiter): Unit = {
+ // For each batch.
+ val batchWriteItemSpec = new BatchWriteItemSpec().withReturnConsumedCapacity(ReturnConsumedCapacity.TOTAL)
+
+ val tableWriteItems = new TableWriteItems(tableName)
+ val tableWriteItemsWithItems: TableWriteItems =
+ // check if hash key only or also range key
+ columnSchema.keys() match {
+ case Left((hashKey, hashKeyIndex, hashKeyType)) =>
+ val hashKeys = items.map(row =>
+ JavaConverter.convertRowValue(row, hashKeyIndex, hashKeyType).asInstanceOf[AnyRef])
+ tableWriteItems.withHashOnlyKeysToDelete(hashKey, hashKeys: _*)
+
+ case Right(((hashKey, hashKeyIndex, hashKeyType), (rangeKey, rangeKeyIndex, rangeKeyType))) =>
+ val alternatingHashAndRangeKeys = items.flatMap { case row =>
+ val hashKeyValue = JavaConverter.convertRowValue(row, hashKeyIndex, hashKeyType)
+ val rangeKeyValue = JavaConverter.convertRowValue(row, rangeKeyIndex, rangeKeyType)
+ Seq(hashKeyValue.asInstanceOf[AnyRef], rangeKeyValue.asInstanceOf[AnyRef])
+ }
+ tableWriteItems.withHashAndRangeKeysToDelete(hashKey, rangeKey, alternatingHashAndRangeKeys: _*)
+ }
+
+ batchWriteItemSpec.withTableWriteItems(tableWriteItemsWithItems)
+
+ val response = client.batchWriteItem(batchWriteItemSpec)
+ handleBatchWriteResponse(client, rateLimiter)(response)
+ }
+
  @tailrec
  private def handleBatchWriteResponse(client: DynamoDB, rateLimiter: RateLimiter)
  (response: BatchWriteItemOutcome): Unit = {

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

@@ -0,0 +1,18 @@
+package com.audienceproject.spark.dynamodb.datasource
+
+import com.amazonaws.services.dynamodbv2.document.DynamoDB
+import com.audienceproject.spark.dynamodb.connector.{ColumnSchema, TableConnector}
+
+class DynamoBatchDeleter(batchSize: Int,
+ columnSchema: ColumnSchema,
+ connector: TableConnector,
+ client: DynamoDB)
+ extends DynamoBatchWriter(batchSize, columnSchema, connector, client) {
+
+ protected override def flush(): Unit = {
+ if (buffer.nonEmpty) {
+ connector.deleteItems(columnSchema, buffer)(client, rateLimiter)
+ buffer.clear()
+ }
+ }
+}

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

@@ -31,13 +31,12 @@ import scala.collection.mutable.ArrayBuffer
 class DynamoBatchWriter(batchSize: Int,
  columnSchema: ColumnSchema,
  connector: TableConnector,
- client: DynamoDB,
- delete: Boolean
+ client: DynamoDB
  )
  extends DataWriter[InternalRow] {
 
- private val buffer = new ArrayBuffer[InternalRow](batchSize)
- private val rateLimiter = RateLimiter.create(connector.writeLimit)
+ protected val buffer = new ArrayBuffer[InternalRow](batchSize)
+ protected val rateLimiter = RateLimiter.create(connector.writeLimit)
 
  override def write(record: InternalRow): Unit = {
  buffer += record.copy()
@@ -53,9 +52,9 @@ class DynamoBatchWriter(batchSize: Int,
 
  override def abort(): Unit = {}
 
- private def flush(): Unit = {
+ protected def flush(): Unit = {
  if (buffer.nonEmpty) {
- connector.putItems(columnSchema, buffer)(client, rateLimiter, delete)
+ connector.putItems(columnSchema, buffer)(client, rateLimiter)
  buffer.clear()
  }
  }

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

@@ -44,8 +44,11 @@ class DynamoWriterFactory(connector: TableConnector,
  assert(!delete, "Please provide exactly one of 'update' or 'delete' options.")
  new DynamoUpdateWriter(columnSchema, connector, client)
  }
+ else if (delete) {
+ new DynamoBatchDeleter(batchSize, columnSchema, connector, client)
+ }
  else
- new DynamoBatchWriter(batchSize, columnSchema, connector, client, delete)
+ new DynamoBatchWriter(batchSize, columnSchema, connector, client)
  }
 
 }

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

@@ -73,7 +73,11 @@ class WriteRelationTest extends AbstractInMemoryTest with Matchers {
  ).toDF("name", "color", "weight")
  newItemsDs.write.dynamodb(tablename)
 
- val toDelete = newItemsDs.filter("name in ('lemon','orange')")
+ val toDelete = Seq(
+ ("lemon", "yellow"),
+ ("orange", "blue"),
+ ("doesn't exist", "black")
+ ).toDF("name", "color")
  toDelete.write.option("delete", "true").dynamodb(tablename)
 
  val validationDs = spark.read.dynamodb(tablename)
@@ -110,10 +114,15 @@ class WriteRelationTest extends AbstractInMemoryTest with Matchers {
  ).toDF("name", "color", "weight")
  newItemsDs.write.dynamodb(tablename)
 
- val toDelete = newItemsDs.filter("color in ('yellow','orange')")
+ val toDelete = Seq(
+ ("lemon", "yellow", 0.1),
+ ("orange", "orange", 0.2),
+ ("pomegranate", "shouldn'tdelete", 0.5)
+ ).toDF("name", "color", "weight")
  toDelete.write.option("delete", "true").dynamodb(tablename)
 
  val validationDs = spark.read.dynamodb(tablename)
+ validationDs.show
  validationDs.count() shouldEqual 2
  validationDs.select("name").as[String].collect should contain theSameElementsAs Seq("lemon", "pomegranate")
  validationDs.select("color").as[String].collect should contain theSameElementsAs Seq("blue", "red")
@@ -143,6 +152,7 @@ class WriteRelationTest extends AbstractInMemoryTest with Matchers {
  .write.option("update", "true").dynamodb(tablename)
 
  val validationDs = spark.read.dynamodb(tablename)
+ validationDs.show
  assert(validationDs.count() === 3)
  assert(validationDs.select("name").as[String].collect().forall(Seq("lemon", "orange", "pomegranate") contains _))
  assert(validationDs.select("color").as[String].collect().forall(Seq("yellow", "orange", "red") contains _))
@@ -169,9 +179,11 @@ class WriteRelationTest extends AbstractInMemoryTest with Matchers {
 
  val alteredDs = newItemsDs
  .withColumn("weight", when($"weight" < 0.2, $"weight").otherwise(lit(null)))
+ alteredDs.show
  alteredDs.write.option("update", "true").dynamodb(tablename)
 
  val validationDs = spark.read.dynamodb(tablename)
+ validationDs.show
  assert(validationDs.count() === 3)
  assert(validationDs.select("name").as[String].collect().forall(Seq("lemon", "orange", "pomegranate") contains _))
  assert(validationDs.select("color").as[String].collect().forall(Seq("yellow", "orange", "red") contains _))