Improved partitioning calculation and added partition reporting

Author: jacobfi

build.sbt

@@ -2,12 +2,14 @@ organization := "com.audienceproject"
 
 name := "spark-dynamodb"
 
-version := "0.5.0"
+version := "1.0.0"
 
 description := "Plug-and-play implementation of an Apache Spark custom data source for AWS DynamoDB."
 
 scalaVersion := "2.11.12"
 
+crossScalaVersions := Seq("2.11.12", "2.12.7")
+
 resolvers += "DynamoDBLocal" at "https://s3-us-west-2.amazonaws.com/dynamodb-local/release"
 
 libraryDependencies += "com.amazonaws" % "aws-java-sdk-sts" % "1.11.678"

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

@@ -60,9 +60,12 @@ private[dynamodb] class TableConnector(tableName: String, parallelism: Int, para
         val itemCount = desc.getItemCount
 
         // Partitioning calculation.
-        val numPartitions = parameters.get("readpartitions").map(_.toInt).getOrElse(
-            (tableSize / maxPartitionBytes).toInt max 1
-        )
+        val numPartitions = parameters.get("readpartitions").map(_.toInt).getOrElse({
+            val sizeBased = (tableSize / maxPartitionBytes).toInt max 1
+            val remainder = sizeBased % parallelism
+            if (remainder > 0) sizeBased + (parallelism - remainder)
+            else sizeBased
+        })
 
         // Provisioned or on-demand throughput.
         val readThroughput = parameters.getOrElse("throughput", Option(desc.getProvisionedThroughput.getReadCapacityUnits)

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

@@ -56,9 +56,12 @@ private[dynamodb] class TableIndexConnector(tableName: String, indexName: String
         val itemCount = indexDesc.getItemCount
 
         // Partitioning calculation.
-        val numPartitions = parameters.get("readpartitions").map(_.toInt).getOrElse(
-            (indexSize / maxPartitionBytes).toInt max 1
-        )
+        val numPartitions = parameters.get("readpartitions").map(_.toInt).getOrElse({
+            val sizeBased = (indexSize / maxPartitionBytes).toInt max 1
+            val remainder = sizeBased % parallelism
+            if (remainder > 0) sizeBased + (parallelism - remainder)
+            else sizeBased
+        })
 
         // Provisioned or on-demand throughput.
         val readThroughput = parameters.getOrElse("throughput", Option(indexDesc.getProvisionedThroughput.getReadCapacityUnits)

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

@@ -26,6 +26,7 @@ import com.audienceproject.spark.dynamodb.connector.{FilterPushdown, TableConnec
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.sources.Filter
 import org.apache.spark.sql.sources.v2.reader._
+import org.apache.spark.sql.sources.v2.reader.partitioning.Partitioning
 import org.apache.spark.sql.types._
 
 import scala.collection.JavaConverters._
@@ -35,7 +36,8 @@ class DynamoDataSourceReader(parallelism: Int,
                              userSchema: Option[StructType] = None)
     extends DataSourceReader
         with SupportsPushDownRequiredColumns
-        with SupportsPushDownFilters {
+        with SupportsPushDownFilters
+        with SupportsReportPartitioning {
 
     private val tableName = parameters("tablename")
     private val indexName = parameters.get("indexName")
@@ -47,6 +49,8 @@ class DynamoDataSourceReader(parallelism: Int,
     private var acceptedFilters: Array[Filter] = Array.empty
     private var currentSchema: StructType = _
 
+    override val outputPartitioning: Partitioning = new OutputPartitioning(dynamoConnector.totalSegments)
+
     override def readSchema(): StructType = {
         if (currentSchema == null)
             currentSchema = userSchema.getOrElse(inferSchema())

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

@@ -0,0 +1,9 @@
+package com.audienceproject.spark.dynamodb.datasource
+
+import org.apache.spark.sql.sources.v2.reader.partitioning.{Distribution, Partitioning}
+
+class OutputPartitioning(override val numPartitions: Int) extends Partitioning {
+
+    override def satisfy(distribution: Distribution): Boolean = false
+
+}