Alias and Doc support in Spark Schema

README.md

@@ -128,6 +128,14 @@ To save DataFrame as avro you should use the `save` method in `AvroSaver`. For e
 ```scala
 scala> AvroSaver.save(myRDD, "my/output/dir")
 ```
+
+To include aliases column in scheme invoke the method `addAvroAliasColumns()` of DataFrame.
+With alias while saving `saveAsAvroFile`, alias columns in DataFrame schema will not be include in Avro file.
+Alias will be available in `aliases` of avro schema.
+```scala
+scala>val dfWithAlias = df.addAvroAliasColumns()
+```
+
 You can also specifiy the the record name and namespace with optional parameters:
 ```scala
 scala> AvroSaver.save(myRDD, "my/output/dir", Map("recordName" -> "MyRecord", "recordNamespace" -> "com.mycompany.mystuff"))

src/main/scala/com/databricks/spark/avro/AvroSaver.scala

@@ -158,7 +158,12 @@ object AvroSaver {
 
             while (convertersIterator.hasNext) {
               val converter = convertersIterator.next()
-              record.put(fieldNamesIterator.next(), converter(rowIterator.next()))
+              val fieldName = fieldNamesIterator.next()
+              if (schema.getField(fieldName) != null) {
+                record.put(fieldName, converter(rowIterator.next()))
+              } else {
+                rowIterator.next()
+              }
             }
             record
           }

src/main/scala/com/databricks/spark/avro/SchemaConverters.scala

@@ -16,19 +16,22 @@
 package com.databricks.spark.avro
 
 import scala.collection.JavaConversions._
-
 import org.apache.avro.{Schema, SchemaBuilder}
 import org.apache.avro.SchemaBuilder._
-
 import org.apache.spark.sql.types._
 import org.apache.avro.Schema.Type._
+import org.apache.spark.sql.DataFrame
+import util.control.Breaks._
 
 /**
  * This object contains method that are used to convert sparkSQL schemas to avro schemas and vice
  * versa.
  */
 private object SchemaConverters {
 
+  val METADATA_KEY_DOC = "doc";
+  val METADATA_KEY_ALIASES = "aliases";
+
   case class SchemaType(dataType: DataType, nullable: Boolean)
 
   /**
@@ -49,7 +52,15 @@ private object SchemaConverters {
       case RECORD =>
         val fields = avroSchema.getFields.map { f =>
           val schemaType = toSqlType(f.schema())
-          StructField(f.name, schemaType.dataType, schemaType.nullable)
+          var meta = new MetadataBuilder()
+          if (f.doc != null) meta.putString(METADATA_KEY_DOC, f.doc)
+          if (f.aliases() != null && f.aliases().size() > 0) {
+            val aliasArray = new Array[String](f.aliases().size())
+            meta.putString("_parent", f.name)
+            f.aliases copyToArray(aliasArray)
+            meta.putStringArray(METADATA_KEY_ALIASES, aliasArray);
+          }
+          StructField(f.name, schemaType.dataType, schemaType.nullable, meta.build())
         }
 
         SchemaType(StructType(fields), nullable = false)
@@ -88,6 +99,19 @@ private object SchemaConverters {
     }
   }
 
+  def dataFrameWithAliasColumn(df : DataFrame) : DataFrame = {
+    var newDf = df
+    for (field <- df.schema.fields) {
+      if (field.metadata.contains(METADATA_KEY_ALIASES)) {
+        val aliasArray = field.metadata.getStringArray(METADATA_KEY_ALIASES)
+        for (alias <- aliasArray) {
+          newDf = newDf.withColumn(alias, df.col(field.name))
+        }
+      }
+    }
+    newDf
+  }
+
   /**
    * This function converts sparkSQL StructType into avro schema. This method uses two other
    * converter methods in order to do the conversion.
@@ -98,14 +122,26 @@ private object SchemaConverters {
       recordNamespace: String): T = {
     val fieldsAssembler: FieldAssembler[T] = schemaBuilder.fields()
     structType.fields.foreach { field =>
-      val newField = fieldsAssembler.name(field.name).`type`()
-
-      if (field.nullable) {
-        convertFieldTypeToAvro(field.dataType, newField.nullable(), field.name, recordNamespace)
-          .noDefault
-      } else {
-        convertFieldTypeToAvro(field.dataType, newField, field.name, recordNamespace)
-          .noDefault
+      breakable {
+        if (field.metadata.contains(METADATA_KEY_ALIASES) && field.metadata.contains("_parent")
+            && !field.metadata.getString("_parent").equals(field.name)) {
+          break
+        }
+        var newFieldBuilder = fieldsAssembler.name(field.name)
+        if (field.metadata contains (METADATA_KEY_DOC)) {
+          newFieldBuilder = newFieldBuilder.doc(field.metadata.getString(METADATA_KEY_DOC))
+        }
+        if (field.metadata.contains(METADATA_KEY_ALIASES)){
+          newFieldBuilder = newFieldBuilder.aliases(field.metadata.getStringArray(METADATA_KEY_ALIASES): _*)
+        }
+        val newField = newFieldBuilder.`type`()
+        if (field.nullable) {
+          convertFieldTypeToAvro(field.dataType, newField.nullable(), field.name, recordNamespace)
+            .noDefault
+        } else {
+          convertFieldTypeToAvro(field.dataType, newField, field.name, recordNamespace)
+            .noDefault
+        }
       }
     }
     fieldsAssembler.endRecord()

src/main/scala/com/databricks/spark/avro/package.scala

@@ -25,6 +25,7 @@ package object avro {
   implicit class AvroContext(sqlContext: SQLContext) {
     def avroFile(filePath: String, minPartitions: Int = 0) =
       sqlContext.baseRelationToDataFrame(AvroRelation(filePath, None, minPartitions)(sqlContext))
+
   }
 
   /**
@@ -35,5 +36,8 @@ package object avro {
         path: String,
         parameters: Map[String, String] = AvroSaver.defaultParameters): Unit =
       AvroSaver.save(dataFrame, path, parameters)
+
+    def addAvroAliasColumns() : DataFrame =
+      SchemaConverters.dataFrameWithAliasColumn(dataFrame)
   }
 }

src/test/resources/test.avsc

@@ -4,25 +4,32 @@
   "fields" : [{
     "name" : "string",
     "type" : "string",
-    "doc"  : "Meaningless string of characters"
+    "doc"  : "Meaningless string of characters",
+    "aliases" : ["string_alias1", "string_alias2"]
   }, {
     "name" : "simple_map",
-    "type" : {"type": "map", "values": "int"}
+    "type" : {"type": "map", "values": "int"},
+    "aliases" : ["map_alias"]
   }, {
     "name" : "complex_map",
-    "type" : {"type": "map", "values": {"type": "map", "values": "string"}}
+    "type" : {"type": "map", "values": {"type": "map", "values": "string"}},
+    "aliases" : ["complex_map_alias"]
   }, {
     "name" : "union_string_null",
-    "type" : ["null", "string"]
+    "type" : ["null", "string"],
+    "aliases" : ["union_string_alias"]
   }, {
     "name" : "union_int_long_null",
-    "type" : ["int", "long", "null"]
+    "type" : ["int", "long", "null"],
+    "aliases" : ["union_int_alias"]
   }, {
     "name" : "union_float_double",
-    "type" : ["float", "double"]
+    "type" : ["float", "double"],
+    "aliases" : ["union_float_alias1", "union_float_alias2"]
   }, {
     "name": "fixed3",
-    "type": {"type": "fixed", "size": 3, "name": "fixed3"}
+    "type": {"type": "fixed", "size": 3, "name": "fixed3"},
+    "aliases" : ["fixed3_alias"]
   }, {
     "name": "fixed2",
     "type": {"type": "fixed", "size": 2, "name": "fixed2"}
@@ -31,7 +38,8 @@
     "type": { "type": "enum",
               "name": "Suit",
               "symbols" : ["SPADES", "HEARTS", "DIAMONDS", "CLUBS"]
-            }
+            },
+    "aliases" : ["enum_alias"]
   }, {
     "name": "record",
     "type": {
@@ -40,14 +48,17 @@
       "aliases": ["RecordAlias"],
       "fields" : [{
         "name": "value_field",
-        "type": "string"
+        "type": "string",
+        "aliases" : ["value_field_alias"]
       }]
     }
   }, {
     "name": "array_of_boolean",
-    "type": {"type": "array", "items": "boolean"}
+    "type": {"type": "array", "items": "boolean"},
+    "aliases" : ["array_of_boolean_alias"]
   }, {
     "name": "bytes",
-    "type": "bytes"
+    "type": "bytes",
+    "aliases" : ["bytes_alias"]
   }]
 }

src/test/scala/com/databricks/spark/avro/AvroSuite.scala

@@ -394,4 +394,93 @@ class AvroSuite extends FunSuite {
     assert(newDf.count == 8)
   }
 
+  test("test doc in meta") {
+    val df = TestSQLContext.load(episodesFile, "com.databricks.spark.avro")
+    df.schema.fields(0).metadata.getString(SchemaConverters.METADATA_KEY_DOC)
+
+    for (x <- df.schema.fields) {
+      if (x.name == "title") {
+        assert("episode title" == x.metadata.getString(SchemaConverters.METADATA_KEY_DOC))
+      } else if (x.name == "doctor") {
+        assert("main actor playing the Doctor in episode" ==
+          x.metadata.getString(SchemaConverters.METADATA_KEY_DOC))
+      } else if (x.name == "air_date") {
+        assert("initial date" == x.metadata.getString(SchemaConverters.METADATA_KEY_DOC))
+      }
+    }
+  }
+
+  test("test aliases in meta") {
+    val df = TestSQLContext.load(testFile, "com.databricks.spark.avro")
+
+    for (x <- df.schema.fields) {
+      if (x.name == "string") {
+        assert(x.metadata.contains(SchemaConverters.METADATA_KEY_ALIASES))
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES).size == 2)
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES)(0) == "string_alias1")
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES)(1) == "string_alias2")
+      } else if (x.name == "simple_map") {
+        assert(x.metadata.contains(SchemaConverters.METADATA_KEY_ALIASES))
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES).size == 1)
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES)(0) == "map_alias")
+      } else if (x.name == "complex_map") {
+        assert(x.metadata.contains(SchemaConverters.METADATA_KEY_ALIASES))
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES).size == 1)
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES)(0) == "complex_map_alias")
+      } else if (x.name == "union_string_null") {
+        assert(x.metadata.contains(SchemaConverters.METADATA_KEY_ALIASES))
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES).size == 1)
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES)(0) == "union_string_alias")
+      } else if (x.name == "union_int_long_null") {
+        assert(x.metadata.contains(SchemaConverters.METADATA_KEY_ALIASES))
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES).size == 1)
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES)(0) == "union_int_alias")
+      } else if (x.name == "union_float_double") {
+        assert(x.metadata.contains(SchemaConverters.METADATA_KEY_ALIASES))
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES).size == 2)
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES)(0) == "union_float_alias1")
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES)(1) == "union_float_alias2")
+      } else if (x.name == "fixed3") {
+        assert(x.metadata.contains(SchemaConverters.METADATA_KEY_ALIASES))
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES).size == 1)
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES)(0) == "fixed3_alias")
+      } else if (x.name == "enum") {
+        assert(x.metadata.contains(SchemaConverters.METADATA_KEY_ALIASES))
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES).size == 1)
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES)(0) == "enum_alias")
+      } else if (x.name == "value_field") {
+        assert(x.metadata.contains(SchemaConverters.METADATA_KEY_ALIASES))
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES).size == 1)
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES)(0) == "value_field_alias")
+      } else if (x.name == "array_of_boolean") {
+        assert(x.metadata.contains(SchemaConverters.METADATA_KEY_ALIASES))
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES).size == 1)
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES)(0) == "array_of_boolean_alias")
+      } else if (x.name == "bytes") {
+        assert(x.metadata.contains(SchemaConverters.METADATA_KEY_ALIASES))
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES).size == 1)
+        assert(x.metadata.getStringArray(SchemaConverters.METADATA_KEY_ALIASES)(0) == "bytes_alias")
+      }
+    }
+  }
+
+  test("test aliases columns in data frame") {
+    var df = TestSQLContext.load(testFile, "com.databricks.spark.avro")
+    var fieldArray = df.schema.fieldNames;
+    assert(fieldArray contains("string"))
+    assert(!(fieldArray contains("string_alias1")))
+    assert(!(fieldArray contains("string_alias2")))
+    assert(!(fieldArray contains("map_alias")))
+    assert(!(fieldArray contains("enum_alias")))
+    assert(!(fieldArray contains("union_int_alias")))
+
+    fieldArray = SchemaConverters.dataFrameWithAliasColumn(df).schema.fieldNames
+    assert(fieldArray contains("string"))
+    assert(fieldArray contains("string_alias1"))
+    assert(fieldArray contains("string_alias2"))
+    assert(fieldArray contains("map_alias"))
+    assert(fieldArray contains("enum_alias"))
+    assert(fieldArray contains("union_int_alias"))
+  }
+
 }