Implemented filter pushdown and column pruning. Added data source to META-INF register

Author: jacobfi

src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister

@@ -0,0 +1 @@
+com.audienceproject.spark.dynamodb.datasource.DefaultSource

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

@@ -104,6 +104,8 @@ private[dynamodb] trait DynamoConnector {
 
     val totalSegments: Int
 
+    val filterPushdownEnabled: Boolean
+
     def scan(segmentNum: Int, columns: Seq[String], filters: Seq[Filter]): ItemCollection[ScanOutcome]
 
     def isEmpty: Boolean = itemLimit == 0

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

@@ -20,7 +20,7 @@
   */
 package com.audienceproject.spark.dynamodb.connector
 
-import com.amazonaws.services.dynamodbv2.xspec.ExpressionSpecBuilder.{N => newN, S => newS, BOOL => newBOOL, _}
+import com.amazonaws.services.dynamodbv2.xspec.ExpressionSpecBuilder.{BOOL => newBOOL, N => newN, S => newS, _}
 import com.amazonaws.services.dynamodbv2.xspec._
 import org.apache.spark.sql.sources._
 
@@ -29,6 +29,23 @@ private[dynamodb] object FilterPushdown {
     def apply(filters: Seq[Filter]): Condition =
         filters.map(buildCondition).map(parenthesize).reduce[Condition](_ and _)
 
+    /**
+      * Accepts only filters that would be considered valid input to FilterPushdown.apply()
+      *
+      * @param filters input list which may contain both valid and invalid filters
+      * @return a (valid, invalid) partitioning of the input filters
+      */
+    def acceptFilters(filters: Array[Filter]): (Array[Filter], Array[Filter]) =
+        filters.partition(checkFilter)
+
+    private def checkFilter(filter: Filter): Boolean = filter match {
+        case _: StringEndsWith => false
+        case And(left, right) => checkFilter(left) && checkFilter(right)
+        case Or(left, right) => checkFilter(left) && checkFilter(right)
+        case Not(f) => checkFilter(f)
+        case _ => true
+    }
+
     private def buildCondition(filter: Filter): Condition = filter match {
         case EqualTo(path, value: Boolean) => newBOOL(path).eq(value)
         case EqualTo(path, value) => coerceAndApply(_ eq _, _ eq _)(path, value)

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

@@ -42,6 +42,8 @@ private[dynamodb] class TableConnector(tableName: String, parallelism: Int, para
     private val region = parameters.get("region")
     private val roleArn = parameters.get("rolearn")
 
+    override val filterPushdownEnabled: Boolean = filterPushdown
+
     override val (keySchema, readLimit, writeLimit, itemLimit, totalSegments) = {
         val table = getDynamoDB(region, roleArn).getTable(tableName)
         val desc = table.describe()

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

@@ -36,6 +36,8 @@ private[dynamodb] class TableIndexConnector(tableName: String, indexName: String
     private val region = parameters.get("region")
     private val roleArn = parameters.get("roleArn")
 
+    override val filterPushdownEnabled: Boolean = filterPushdown
+
     override val (keySchema, readLimit, itemLimit, totalSegments) = {
         val table = getDynamoDB(region, roleArn).getTable(tableName)
         val indexDesc = table.describe().getGlobalSecondaryIndexes.asScala.find(_.getIndexName == indexName).get

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

@@ -36,6 +36,12 @@ class DefaultSource extends ReadSupport with WriteSupport with DataSourceRegiste
 
     private val logger = LoggerFactory.getLogger(this.getClass)
 
+    override def createReader(schema: StructType, options: DataSourceOptions): DataSourceReader = {
+        val optionsMap = options.asMap().asScala
+        val defaultParallelism = optionsMap.get("defaultparallelism").map(_.toInt).getOrElse(getDefaultParallelism)
+        new DynamoDataSourceReader(defaultParallelism, Map(optionsMap.toSeq: _*), Some(schema))
+    }
+
     override def createReader(options: DataSourceOptions): DataSourceReader = {
         val optionsMap = options.asMap().asScala
         val defaultParallelism = optionsMap.get("defaultparallelism").map(_.toInt).getOrElse(getDefaultParallelism)

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

@@ -22,17 +22,20 @@ package com.audienceproject.spark.dynamodb.datasource
 
 import java.util
 
-import com.audienceproject.spark.dynamodb.connector.{TableConnector, TableIndexConnector}
+import com.audienceproject.spark.dynamodb.connector.{FilterPushdown, TableConnector, TableIndexConnector}
 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.types._
 
 import scala.collection.JavaConverters._
 
-class DynamoDataSourceReader(parallelism: Int, parameters: Map[String, String]) extends DataSourceReader
-    with SupportsPushDownRequiredColumns
-    with SupportsPushDownFilters {
+class DynamoDataSourceReader(parallelism: Int,
+                             parameters: Map[String, String],
+                             userSchema: Option[StructType] = None)
+    extends DataSourceReader
+        with SupportsPushDownRequiredColumns
+        with SupportsPushDownFilters {
 
     private val tableName = parameters("tablename")
     private val indexName = parameters.get("indexName")
@@ -41,21 +44,41 @@ class DynamoDataSourceReader(parallelism: Int, parameters: Map[String, String])
         if (indexName.isDefined) new TableIndexConnector(tableName, indexName.get, parallelism, parameters)
         else new TableConnector(tableName, parallelism, parameters)
 
-    override lazy val readSchema: StructType = inferSchema()
+    private var acceptedFilters: Array[Filter] = Array.empty
+    private var currentSchema: StructType = _
+
+    override def readSchema(): StructType = {
+        if (currentSchema == null)
+            currentSchema = userSchema.getOrElse(inferSchema())
+        currentSchema
+    }
 
     override def planInputPartitions(): util.List[InputPartition[InternalRow]] = {
-        val inputPartitions = new util.ArrayList[InputPartition[InternalRow]]()
+        val inputPartitions = new util.ArrayList[InputPartition[InternalRow]]
         for (partitionIndex <- 0 until dynamoConnector.totalSegments) {
-            inputPartitions.add(new ScanPartition(readSchema, partitionIndex, dynamoConnector))
+            inputPartitions.add(new ScanPartition(readSchema(), partitionIndex, dynamoConnector, acceptedFilters))
         }
         inputPartitions
     }
 
-    override def pruneColumns(requiredSchema: StructType): Unit = ???
+    override def pruneColumns(requiredSchema: StructType): Unit = {
+        val schema = readSchema()
+        val keyFields = Seq(Some(dynamoConnector.keySchema.hashKeyName), dynamoConnector.keySchema.rangeKeyName).flatten
+            .flatMap(keyName => schema.fields.find(_.name == keyName))
+        val requiredFields = keyFields ++ requiredSchema.fields
+        val newFields = readSchema().fields.filter(requiredFields.contains)
+        currentSchema = StructType(newFields)
+    }
 
-    override def pushFilters(filters: Array[Filter]): Array[Filter] = ???
+    override def pushFilters(filters: Array[Filter]): Array[Filter] = {
+        if (dynamoConnector.filterPushdownEnabled) {
+            val (acceptedFilters, postScanFilters) = FilterPushdown.acceptFilters(filters)
+            this.acceptedFilters = acceptedFilters
+            postScanFilters // Return filters that need to be evaluated after scanning.
+        } else filters
+    }
 
-    override def pushedFilters(): Array[Filter] = ???
+    override def pushedFilters(): Array[Filter] = acceptedFilters
 
     private def inferSchema(): StructType = {
         val inferenceItems =

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

@@ -24,14 +24,16 @@ import com.amazonaws.services.dynamodbv2.document.Item
 import com.audienceproject.spark.dynamodb.connector.DynamoConnector
 import com.audienceproject.spark.dynamodb.util.RateLimiter
 import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.sources.Filter
 import org.apache.spark.sql.sources.v2.reader.{InputPartition, InputPartitionReader}
 import org.apache.spark.sql.types.{StructField, StructType}
 
 import scala.collection.JavaConverters._
 
 class ScanPartition(schema: StructType,
                     partitionIndex: Int,
-                    connector: DynamoConnector)
+                    connector: DynamoConnector,
+                    filters: Array[Filter])
     extends InputPartition[InternalRow] {
 
     private val requiredColumns = schema.map(_.name)
@@ -56,7 +58,7 @@ class ScanPartition(schema: StructType,
 
     private class PartitionReader extends InputPartitionReader[InternalRow] {
 
-        private val pageIterator = connector.scan(partitionIndex, Seq.empty, Seq.empty).pages().iterator().asScala
+        private val pageIterator = connector.scan(partitionIndex, requiredColumns, filters).pages().iterator().asScala
         private val rateLimiter = new RateLimiter(connector.readLimit)
 
         private var innerIterator: Iterator[InternalRow] = Iterator.empty