Fixed issue #76. See PR #87

Author: jacobfi

src/main/scala/com/audienceproject/spark/dynamodb/implicits.scala

@@ -24,7 +24,6 @@ import com.audienceproject.spark.dynamodb.reflect.SchemaAnalysis
 import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.functions.col
-import org.apache.spark.sql.types.StructField
 
 import scala.reflect.ClassTag
 import scala.reflect.runtime.universe.TypeTag
@@ -41,28 +40,29 @@ object implicits {
 
         def dynamodbAs[T <: Product : ClassTag : TypeTag](tableName: String): Dataset[T] = {
             implicit val encoder: Encoder[T] = ExpressionEncoder()
-            getColumnsAlias(getDynamoDBSource(tableName)
-                .schema(SchemaAnalysis[T]).load()).as
+            val (schema, aliasMap) = SchemaAnalysis[T]
+            getColumnsAlias(getDynamoDBSource(tableName).schema(schema).load(), aliasMap).as
         }
 
         def dynamodbAs[T <: Product : ClassTag : TypeTag](tableName: String, indexName: String): Dataset[T] = {
             implicit val encoder: Encoder[T] = ExpressionEncoder()
-            getColumnsAlias(getDynamoDBSource(tableName)
-                .option("indexName", indexName)
-                .schema(SchemaAnalysis[T]).load()).as
+            val (schema, aliasMap) = SchemaAnalysis[T]
+            getColumnsAlias(
+                getDynamoDBSource(tableName).option("indexName", indexName).schema(schema).load(), aliasMap).as
         }
 
         private def getDynamoDBSource(tableName: String): DataFrameReader =
             reader.format("com.audienceproject.spark.dynamodb.datasource").option("tableName", tableName)
 
-        private def getColumnsAlias(dataFrame: DataFrame): DataFrame = {
-            val columnsAlias = dataFrame.schema.collect({
-                case StructField(name, _, _, metadata) if metadata.contains("alias") =>
-                    col(name).as(metadata.getString("alias"))
-                case StructField(name, _, _, _) =>
-                    col(name)
-            })
-            dataFrame.select(columnsAlias: _*)
+        private def getColumnsAlias(dataFrame: DataFrame, aliasMap: Map[String, String]): DataFrame = {
+            if (aliasMap.isEmpty) dataFrame
+            else {
+                val columnsAlias = dataFrame.columns.map({
+                    case name if aliasMap.isDefinedAt(name) => col(name) as aliasMap(name)
+                    case name => col(name)
+                })
+                dataFrame.select(columnsAlias: _*)
+            }
         }
 
     }

src/main/scala/com/audienceproject/spark/dynamodb/reflect/SchemaAnalysis.scala

@@ -22,7 +22,7 @@ package com.audienceproject.spark.dynamodb.reflect
 
 import com.audienceproject.spark.dynamodb.attribute
 import org.apache.spark.sql.catalyst.ScalaReflection
-import org.apache.spark.sql.types.{Metadata, MetadataBuilder, StructField, StructType}
+import org.apache.spark.sql.types.{StructField, StructType}
 
 import scala.reflect.ClassTag
 import scala.reflect.runtime.{universe => ru}
@@ -32,35 +32,37 @@ import scala.reflect.runtime.{universe => ru}
   */
 private[dynamodb] object SchemaAnalysis {
 
-    def apply[T <: Product : ClassTag : ru.TypeTag]: StructType = {
+    def apply[T <: Product : ClassTag : ru.TypeTag]: (StructType, Map[String, String]) = {
 
         val runtimeMirror = ru.runtimeMirror(getClass.getClassLoader)
 
         val classObj = scala.reflect.classTag[T].runtimeClass
         val classSymbol = runtimeMirror.classSymbol(classObj)
 
-        val sparkFields = classSymbol.primaryConstructor.typeSignature.paramLists.head.map(field => {
-            val sparkType = ScalaReflection.schemaFor(field.typeSignature).dataType
+        val params = classSymbol.primaryConstructor.typeSignature.paramLists.head
+        val (sparkFields, aliasMap) = params.foldLeft((List.empty[StructField], Map.empty[String, String]))({
+            case ((list, map), field) =>
+                val sparkType = ScalaReflection.schemaFor(field.typeSignature).dataType
 
-            // Black magic from here:
-            // https://stackoverflow.com/questions/23046958/accessing-an-annotation-value-in-scala
-            val attrName = field.annotations.collectFirst({
-                case ann: ru.AnnotationApi if ann.tree.tpe =:= ru.typeOf[attribute] =>
-                    ann.tree.children.tail.collectFirst({
-                        case ru.Literal(ru.Constant(name: String)) => name
-                    })
-            }).flatten
+                // Black magic from here:
+                // https://stackoverflow.com/questions/23046958/accessing-an-annotation-value-in-scala
+                val attrName = field.annotations.collectFirst({
+                    case ann: ru.AnnotationApi if ann.tree.tpe =:= ru.typeOf[attribute] =>
+                        ann.tree.children.tail.collectFirst({
+                            case ru.Literal(ru.Constant(name: String)) => name
+                        })
+                }).flatten
 
-            if (attrName.isDefined) {
-                val metadata = new MetadataBuilder().putString("alias", field.name.toString).build()
-                StructField(attrName.get, sparkType, nullable = true, metadata)
-            } else {
-                StructField(field.name.toString, sparkType, nullable = true, Metadata.empty)
-            }
+                if (attrName.isDefined) {
+                    val sparkField = StructField(attrName.get, sparkType, nullable = true)
+                    (list :+ sparkField, map + (attrName.get -> field.name.toString))
+                } else {
+                    val sparkField = StructField(field.name.toString, sparkType, nullable = true)
+                    (list :+ sparkField, map)
+                }
         })
 
-        StructType(sparkFields)
-
+        (StructType(sparkFields), aliasMap)
     }
 
 }

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

@@ -57,6 +57,10 @@ class DefaultSourceTest extends AbstractInMemoryTest {
 
     test("Test of attribute name alias") {
         import spark.implicits._
+        spark.read.dynamodb("TestFruit").printSchema()
+        spark.read.dynamodb("TestFruit").show()
+        spark.read.dynamodbAs[TestFruit]("TestFruit").printSchema()
+        spark.read.dynamodbAs[TestFruit]("TestFruit").show()
         val itemApple = spark.read.dynamodbAs[TestFruit]("TestFruit")
             .filter($"primaryKey" === "apple")
             .takeAsList(1).get(0)