Merge pull request #33 from audienceproject/feature/throughput-as-option

Add the option to hardcode read and write throughput

Author: cosmincatalin

README.md

@@ -81,12 +81,14 @@ The following parameters can be set as options on the Spark reader object before
 - `stronglyConsistentReads` whether or not to use strongly consistent reads. Default false.
 - `bytesPerRCU` number of bytes that can be read per second with a single Read Capacity Unit. Default 4000 (4 KB). This value is multiplied by two when `stronglyConsistentReads=false`
 - `filterPushdown` whether or not to use filter pushdown to DynamoDB on scan requests. Default true.
+- `throughput` the desired read throughput to use. It overwrites any calculation used by the package. It is intended to be used with tables that are on-demand. Defaults to 100 for on-demand.
 
 The following parameters can be set as options on the Spark writer object before saving.
 
 - `writeBatchSize` number of items to send per call to DynamoDB BatchWriteItem. Default 25.
 - `targetCapacity` fraction of provisioned write capacity on the table to consume for writing or updating. Default 1 (i.e. 100% capacity).
-- `update` if true items will be written using UpdateItem on keys rather than BatchWriteItem. Default false. 
+- `update` if true items will be written using UpdateItem on keys rather than BatchWriteItem. Default false.
+- `throughput` the desired write throughput to use. It overwrites any calculation used by the package. It is intended to be used with tables that are on-demand. Defaults to 100 for on-demand.
 
 ## Running Unit Tests
 The unit tests are dependent on the AWS DynamoDBLocal client, which in turn is dependent on [sqlite4java](https://bitbucket.org/almworks/sqlite4java/src/master/). I had some problems running this on OSX, so I had to put the library directly in the /lib folder, as graciously explained in [this Stack Overflow answer](https://stackoverflow.com/questions/34137043/amazon-dynamodb-local-unknown-error-exception-or-failure/35353377#35353377).

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

@@ -56,12 +56,12 @@ private[dynamodb] class TableConnector(tableName: String, totalSegments: Int, pa
         val readFactor = if (consistentRead) 1 else 2
 
         // Provisioned or on-demand throughput.
-        val readThroughput = Option(desc.getProvisionedThroughput.getReadCapacityUnits)
-            .filter(_ > 0).map(_.longValue())
-            .getOrElse(100L)
-        val writeThroughput = Option(desc.getProvisionedThroughput.getWriteCapacityUnits)
-            .filter(_ > 0).map(_.longValue())
-            .getOrElse(100L)
+        val readThroughput = parameters.getOrElse("throughput", Option(desc.getProvisionedThroughput.getReadCapacityUnits)
+            .filter(_ > 0).map(_.longValue().toString)
+            .getOrElse("100")).toLong
+        val writeThroughput = parameters.getOrElse("throughput", Option(desc.getProvisionedThroughput.getWriteCapacityUnits)
+            .filter(_ > 0).map(_.longValue().toString)
+            .getOrElse("100")).toLong
 
         // Rate limit calculation.
         val tableSize = desc.getTableSizeBytes

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

@@ -49,9 +49,9 @@ private[dynamodb] class TableIndexConnector(tableName: String, indexName: String
         val readFactor = if (consistentRead) 1 else 2
 
         // Provisioned or on-demand throughput.
-        val readThroughput = Option(indexDesc.getProvisionedThroughput.getReadCapacityUnits)
-            .filter(_ > 0).map(_.longValue())
-            .getOrElse(100L)
+        val readThroughput = parameters.getOrElse("throughput", Option(indexDesc.getProvisionedThroughput.getReadCapacityUnits)
+            .filter(_ > 0).map(_.longValue().toString)
+            .getOrElse("100")).toLong
 
         // Rate limit calculation.
         val tableSize = indexDesc.getIndexSizeBytes