Attempting to fix loss of schema

Author: AndreSchumacher

sql/core/src/main/scala/org/apache/spark/sql/parquet/ParquetRelation.scala

@@ -55,6 +55,8 @@ private[sql] case class ParquetRelation(val path: String)
       .getSchema
 
   /** Attributes */
+  // TODO: THIS POTENTIALLY LOOSES TYPE INFORMATION!!!!
+  // e.g. short <-> INT32 and byte <-> INT32
   override val output =
     ParquetTypesConverter
       .convertToAttributes(parquetSchema)
@@ -132,7 +134,9 @@ private[sql] object ParquetRelation {
     }
     ParquetRelation.enableLogForwarding()
     ParquetTypesConverter.writeMetaData(attributes, path, conf)
-    new ParquetRelation(path.toString)
+    new ParquetRelation(path.toString) {
+      override val output = attributes
+    }
   }
 
   private def checkPath(pathStr: String, allowExisting: Boolean, conf: Configuration): Path = {

sql/core/src/main/scala/org/apache/spark/sql/parquet/ParquetTableSupport.scala

@@ -126,12 +126,7 @@ private[parquet] class RowWriteSupport extends WriteSupport[Row] with Logging {
 
   override def init(configuration: Configuration): WriteSupport.WriteContext = {
     //attributes = DataType(configuration.get(RowWriteSupport.PARQUET_ROW_SCHEMA))
-    attributes = if (attributes == null) {
-      RowWriteSupport.getSchema(configuration) match {
-        case s: StructType => s.toAttributes
-        case other => sys.error(s"Can convert $attributes to row")
-      }
-    } else attributes
+    attributes = if (attributes == null) RowWriteSupport.getSchema(configuration) else attributes
     schema = if (schema == null) ParquetTypesConverter.convertFromAttributes(attributes) else schema
     // ParquetTypesConverter.convertToAttributes(schema)
     log.debug(s"write support initialized for requested schema $attributes")
@@ -192,7 +187,9 @@ private[parquet] class RowWriteSupport extends WriteSupport[Row] with Logging {
           )
         )
         case IntegerType => writer.addInteger(value.asInstanceOf[Int])
+        case ShortType => writer.addInteger(value.asInstanceOf[Int])
         case LongType => writer.addLong(value.asInstanceOf[Long])
+        case ByteType => writer.addInteger(value.asInstanceOf[Int])
         case DoubleType => writer.addDouble(value.asInstanceOf[Double])
         case FloatType => writer.addFloat(value.asInstanceOf[Float])
         case BooleanType => writer.addBoolean(value.asInstanceOf[Boolean])
@@ -298,7 +295,9 @@ private[parquet] class MutableRowWriteSupport extends RowWriteSupport {
         )
       )
       case IntegerType => writer.addInteger(record.getInt(index))
+      case ShortType => writer.addInteger(record.getShort(index))
       case LongType => writer.addLong(record.getLong(index))
+      case ByteType => writer.addInteger(record.getByte(index))
       case DoubleType => writer.addDouble(record.getDouble(index))
       case FloatType => writer.addFloat(record.getFloat(index))
       case BooleanType => writer.addBoolean(record.getBoolean(index))

sql/core/src/main/scala/org/apache/spark/sql/parquet/ParquetTypes.scala

@@ -292,13 +292,14 @@ private[parquet] object ParquetTypesConverter {
   }
 
   def convertFromString(string: String): Seq[Attribute] = {
-    val decoded: Array[Byte] = BaseEncoding.base64().decode(string)
-    SparkSqlSerializer.deserialize(decoded)
+    DataType(string) match {
+      case s: StructType => s.toAttributes
+      case other => sys.error(s"Can convert $string to row")
+    }
   }
 
   def convertToString(schema: Seq[Attribute]): String = {
-    val serialized: Array[Byte] = SparkSqlSerializer.serialize(schema)
-    BaseEncoding.base64().encode(serialized)
+    StructType.fromAttributes(schema).toString
   }
 
   def writeMetaData(attributes: Seq[Attribute], origPath: Path, conf: Configuration) {