Upgraded to Scala 2.12.12 and Spark 3.0.0

Author: jacobfi

build.sbt

@@ -2,13 +2,11 @@ organization := "com.audienceproject"
 
 name := "spark-dynamodb"
 
-version := "1.0.5"
+version := "1.1.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")
+scalaVersion := "2.12.12"
 
 compileOrder := CompileOrder.JavaThenScala
 
@@ -18,7 +16,7 @@ libraryDependencies += "com.amazonaws" % "aws-java-sdk-sts" % "1.11.678"
 libraryDependencies += "com.amazonaws" % "aws-java-sdk-dynamodb" % "1.11.678"
 libraryDependencies += "com.amazonaws" % "DynamoDBLocal" % "[1.11,2.0)" % "test" exclude("com.google.guava", "guava")
 
-libraryDependencies += "org.apache.spark" %% "spark-sql" % "2.4.4" % "provided"
+libraryDependencies += "org.apache.spark" %% "spark-sql" % "3.0.0" % "provided"
 
 libraryDependencies += "org.scalatest" %% "scalatest" % "3.0.5" % "test"
 
@@ -53,7 +51,7 @@ Test / resourceGenerators += Def.task {
     import java.util.stream.Collectors
     import scala.collection.JavaConverters._
 
-    def log(msg: Any) = println(s"[℣₳ℒ𐎅] $msg") //stand out in the crowd
+    def log(msg: Any): Unit = println(s"[℣₳ℒ𐎅] $msg") //stand out in the crowd
 
     val theOnesWeLookFor = Set(
         "libsqlite4java-linux-amd64-1.0.392.so",

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

@@ -20,53 +20,28 @@
   */
 package com.audienceproject.spark.dynamodb.datasource
 
-import java.util.Optional
+import java.util
 
+import org.apache.spark.sql.connector.catalog.{Table, TableProvider}
+import org.apache.spark.sql.connector.expressions.Transform
 import org.apache.spark.sql.sources.DataSourceRegister
-import org.apache.spark.sql.sources.v2.reader.DataSourceReader
-import org.apache.spark.sql.sources.v2.writer.DataSourceWriter
-import org.apache.spark.sql.sources.v2.{DataSourceOptions, ReadSupport, WriteSupport}
 import org.apache.spark.sql.types.StructType
-import org.apache.spark.sql.{SaveMode, SparkSession}
-import org.slf4j.LoggerFactory
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
 
-import scala.collection.JavaConverters._
+class DefaultSource extends TableProvider with DataSourceRegister {
 
-class DefaultSource extends ReadSupport with WriteSupport with DataSourceRegister {
-
-    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 getTable(schema: StructType,
+                          partitioning: Array[Transform],
+                          properties: util.Map[String, String]): Table = {
+        new DynamoTable(new CaseInsensitiveStringMap(properties), Some(schema))
     }
 
-    override def createReader(options: DataSourceOptions): DataSourceReader = {
-        val optionsMap = options.asMap().asScala
-        val defaultParallelism = optionsMap.get("defaultparallelism").map(_.toInt).getOrElse(getDefaultParallelism)
-        new DynamoDataSourceReader(defaultParallelism, Map(optionsMap.toSeq: _*))
+    override def inferSchema(options: CaseInsensitiveStringMap): StructType = {
+        new DynamoTable(options).schema()
     }
 
-    override def createWriter(writeUUID: String, schema: StructType, mode: SaveMode, options: DataSourceOptions): Optional[DataSourceWriter] = {
-        if (mode == SaveMode.Append || mode == SaveMode.Overwrite)
-            throw new IllegalArgumentException(s"DynamoDB data source does not support save modes ($mode)." +
-                " Please use option 'update' (true | false) to differentiate between append/overwrite and append/update behavior.")
-        val optionsMap = options.asMap().asScala
-        val defaultParallelism = optionsMap.get("defaultparallelism").map(_.toInt).getOrElse(getDefaultParallelism)
-        val writer = new DynamoDataSourceWriter(defaultParallelism, Map(optionsMap.toSeq: _*), schema)
-        Optional.of(writer)
-    }
+    override def supportsExternalMetadata(): Boolean = true
 
     override def shortName(): String = "dynamodb"
 
-    private def getDefaultParallelism: Int =
-        SparkSession.getActiveSession match {
-            case Some(spark) => spark.sparkContext.defaultParallelism
-            case None =>
-                logger.warn("Unable to read defaultParallelism from SparkSession." +
-                    " Parallelism will be 1 unless overwritten with option `defaultParallelism`")
-                1
-        }
-
 }

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

@@ -0,0 +1,48 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one
+  * or more contributor license agreements.  See the NOTICE file
+  * distributed with this work for additional information
+  * regarding copyright ownership.  The ASF licenses this file
+  * to you under the Apache License, Version 2.0 (the
+  * "License"); you may not use this file except in compliance
+  * with the License.  You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing,
+  * software distributed under the License is distributed on an
+  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+  * KIND, either express or implied.  See the License for the
+  * specific language governing permissions and limitations
+  * under the License.
+  *
+  * Copyright © 2019 AudienceProject. All rights reserved.
+  */
+package com.audienceproject.spark.dynamodb.datasource
+
+import com.audienceproject.spark.dynamodb.connector.DynamoConnector
+import org.apache.spark.sql.connector.read._
+import org.apache.spark.sql.connector.read.partitioning.Partitioning
+import org.apache.spark.sql.sources.Filter
+import org.apache.spark.sql.types.StructType
+
+class DynamoBatchReader(connector: DynamoConnector,
+                        filters: Array[Filter],
+                        schema: StructType)
+    extends Scan with Batch with SupportsReportPartitioning {
+
+    override def readSchema(): StructType = schema
+
+    override def toBatch: Batch = this
+
+    override def planInputPartitions(): Array[InputPartition] = {
+        val requiredColumns = schema.map(_.name)
+        Array.tabulate(connector.totalSegments)(new ScanPartition(_, requiredColumns, filters))
+    }
+
+    override def createReaderFactory(): PartitionReaderFactory =
+        new DynamoReaderFactory(connector, schema)
+
+    override val outputPartitioning: Partitioning = new OutputPartitioning(connector.totalSegments)
+
+}

src/main/scala/com/audienceproject/spark/dynamodb/datasource/DynamoBatchDeleteWriter.scala renamed to src/main/scala/com/audienceproject/spark/dynamodb/datasource/DynamoDataDeleteWriter.scala

@@ -24,11 +24,11 @@ package com.audienceproject.spark.dynamodb.datasource
 import com.amazonaws.services.dynamodbv2.document.DynamoDB
 import com.audienceproject.spark.dynamodb.connector.{ColumnSchema, TableConnector}
 
-class DynamoBatchDeleteWriter(batchSize: Int,
-                              columnSchema: ColumnSchema,
-                              connector: TableConnector,
-                              client: DynamoDB)
-    extends DynamoBatchWriter(batchSize, columnSchema, connector, client) {
+class DynamoDataDeleteWriter(batchSize: Int,
+                             columnSchema: ColumnSchema,
+                             connector: TableConnector,
+                             client: DynamoDB)
+    extends DynamoDataWriter(batchSize, columnSchema, connector, client) {
 
     protected override def flush(): Unit = {
         if (buffer.nonEmpty) {

src/main/scala/com/audienceproject/spark/dynamodb/datasource/DynamoUpdateWriter.scala renamed to src/main/scala/com/audienceproject/spark/dynamodb/datasource/DynamoDataUpdateWriter.scala

@@ -24,11 +24,11 @@ import com.amazonaws.services.dynamodbv2.document.DynamoDB
 import com.audienceproject.shaded.google.common.util.concurrent.RateLimiter
 import com.audienceproject.spark.dynamodb.connector.{ColumnSchema, TableConnector}
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.sources.v2.writer.{DataWriter, WriterCommitMessage}
+import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage}
 
-class DynamoUpdateWriter(columnSchema: ColumnSchema,
-                         connector: TableConnector,
-                         client: DynamoDB)
+class DynamoDataUpdateWriter(columnSchema: ColumnSchema,
+                             connector: TableConnector,
+                             client: DynamoDB)
     extends DataWriter[InternalRow] {
 
     private val rateLimiter = RateLimiter.create(connector.writeLimit)
@@ -41,4 +41,6 @@ class DynamoUpdateWriter(columnSchema: ColumnSchema,
 
     override def abort(): Unit = {}
 
+    override def close(): Unit = client.shutdown()
+
 }

src/main/scala/com/audienceproject/spark/dynamodb/datasource/DynamoBatchWriter.scala renamed to src/main/scala/com/audienceproject/spark/dynamodb/datasource/DynamoDataWriter.scala

@@ -24,14 +24,14 @@ import com.amazonaws.services.dynamodbv2.document.DynamoDB
 import com.audienceproject.shaded.google.common.util.concurrent.RateLimiter
 import com.audienceproject.spark.dynamodb.connector.{ColumnSchema, TableConnector}
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.sources.v2.writer.{DataWriter, WriterCommitMessage}
+import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage}
 
 import scala.collection.mutable.ArrayBuffer
 
-class DynamoBatchWriter(batchSize: Int,
-                        columnSchema: ColumnSchema,
-                        connector: TableConnector,
-                        client: DynamoDB)
+class DynamoDataWriter(batchSize: Int,
+                       columnSchema: ColumnSchema,
+                       connector: TableConnector,
+                       client: DynamoDB)
     extends DataWriter[InternalRow] {
 
     protected val buffer: ArrayBuffer[InternalRow] = new ArrayBuffer[InternalRow](batchSize)
@@ -51,6 +51,8 @@ class DynamoBatchWriter(batchSize: Int,
 
     override def abort(): Unit = {}
 
+    override def close(): Unit = client.shutdown()
+
     protected def flush(): Unit = {
         if (buffer.nonEmpty) {
             connector.putItems(columnSchema, buffer)(client, rateLimiter)

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

@@ -0,0 +1,99 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one
+  * or more contributor license agreements.  See the NOTICE file
+  * distributed with this work for additional information
+  * regarding copyright ownership.  The ASF licenses this file
+  * to you under the Apache License, Version 2.0 (the
+  * "License"); you may not use this file except in compliance
+  * with the License.  You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing,
+  * software distributed under the License is distributed on an
+  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+  * KIND, either express or implied.  See the License for the
+  * specific language governing permissions and limitations
+  * under the License.
+  *
+  * Copyright © 2019 AudienceProject. All rights reserved.
+  */
+package com.audienceproject.spark.dynamodb.datasource
+
+import com.amazonaws.services.dynamodbv2.document.Item
+import com.audienceproject.shaded.google.common.util.concurrent.RateLimiter
+import com.audienceproject.spark.dynamodb.connector.DynamoConnector
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader, PartitionReaderFactory}
+import org.apache.spark.sql.types.{StructField, StructType}
+
+import scala.collection.JavaConverters._
+
+class DynamoReaderFactory(connector: DynamoConnector,
+                          schema: StructType)
+    extends PartitionReaderFactory {
+
+    override def createReader(partition: InputPartition): PartitionReader[InternalRow] = {
+        if (connector.isEmpty) new EmptyReader
+        else new ScanPartitionReader(partition.asInstanceOf[ScanPartition])
+    }
+
+    private class EmptyReader extends PartitionReader[InternalRow] {
+        override def next(): Boolean = false
+
+        override def get(): InternalRow = throw new IllegalStateException("Unable to call get() on empty iterator")
+
+        override def close(): Unit = {}
+    }
+
+    private class ScanPartitionReader(scanPartition: ScanPartition) extends PartitionReader[InternalRow] {
+
+        import scanPartition._
+
+        private val pageIterator = connector.scan(partitionIndex, requiredColumns, filters).pages().iterator().asScala
+        private val rateLimiter = RateLimiter.create(connector.readLimit)
+
+        private var innerIterator: Iterator[InternalRow] = Iterator.empty
+
+        private var currentRow: InternalRow = _
+        private var proceed = false
+
+        private val typeConversions = schema.collect({
+            case StructField(name, dataType, _, _) => name -> TypeConversion(name, dataType)
+        }).toMap
+
+        override def next(): Boolean = {
+            proceed = true
+            innerIterator.hasNext || {
+                if (pageIterator.hasNext) {
+                    nextPage()
+                    next()
+                }
+                else false
+            }
+        }
+
+        override def get(): InternalRow = {
+            if (proceed) {
+                currentRow = innerIterator.next()
+                proceed = false
+            }
+            currentRow
+        }
+
+        override def close(): Unit = {}
+
+        private def nextPage(): Unit = {
+            val page = pageIterator.next()
+            val result = page.getLowLevelResult
+            Option(result.getScanResult.getConsumedCapacity).foreach(cap => rateLimiter.acquire(cap.getCapacityUnits.toInt max 1))
+            innerIterator = result.getItems.iterator().asScala.map(itemToRow(requiredColumns))
+        }
+
+        private def itemToRow(requiredColumns: Seq[String])(item: Item): InternalRow =
+            if (requiredColumns.nonEmpty) InternalRow.fromSeq(requiredColumns.map(columnName => typeConversions(columnName)(item)))
+            else InternalRow.fromSeq(item.asMap().asScala.values.toSeq.map(_.toString))
+
+    }
+
+}

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

@@ -0,0 +1,56 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one
+  * or more contributor license agreements.  See the NOTICE file
+  * distributed with this work for additional information
+  * regarding copyright ownership.  The ASF licenses this file
+  * to you under the Apache License, Version 2.0 (the
+  * "License"); you may not use this file except in compliance
+  * with the License.  You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing,
+  * software distributed under the License is distributed on an
+  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+  * KIND, either express or implied.  See the License for the
+  * specific language governing permissions and limitations
+  * under the License.
+  *
+  * Copyright © 2019 AudienceProject. All rights reserved.
+  */
+package com.audienceproject.spark.dynamodb.datasource
+
+import com.audienceproject.spark.dynamodb.connector.{DynamoConnector, FilterPushdown}
+import org.apache.spark.sql.connector.read._
+import org.apache.spark.sql.sources.Filter
+import org.apache.spark.sql.types._
+
+class DynamoScanBuilder(connector: DynamoConnector, schema: StructType)
+    extends ScanBuilder
+        with SupportsPushDownRequiredColumns
+        with SupportsPushDownFilters {
+
+    private var acceptedFilters: Array[Filter] = Array.empty
+    private var currentSchema: StructType = schema
+
+    override def build(): Scan = new DynamoBatchReader(connector, pushedFilters(), currentSchema)
+
+    override def pruneColumns(requiredSchema: StructType): Unit = {
+        val keyFields = Seq(Some(connector.keySchema.hashKeyName), connector.keySchema.rangeKeyName).flatten
+            .flatMap(keyName => currentSchema.fields.find(_.name == keyName))
+        val requiredFields = keyFields ++ requiredSchema.fields
+        val newFields = currentSchema.fields.filter(requiredFields.contains)
+        currentSchema = StructType(newFields)
+    }
+
+    override def pushFilters(filters: Array[Filter]): Array[Filter] = {
+        if (connector.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] = acceptedFilters
+
+}