Enable inserting array into hive table saved as parquet using datasource.

viirya · viirya · commit 4e3bd5568e64 · 2015-02-24T01:03:30.000+08:00
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/parquet/ParquetTableSupport.scala b/sql/core/src/main/scala/org/apache/spark/sql/parquet/ParquetTableSupport.scala
@@ -28,7 +28,7 @@ import parquet.io.api._
 import parquet.schema.MessageType
 
 import org.apache.spark.Logging
-import org.apache.spark.sql.catalyst.expressions.{Attribute, Row}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, Row}
 import org.apache.spark.sql.types._
 
 /**
@@ -381,7 +381,14 @@ private[parquet] object RowWriteSupport {
   }
 
   def setSchema(schema: Seq[Attribute], configuration: Configuration) {
-    val encoded = ParquetTypesConverter.convertToString(schema)
+    val updatedSchama = schema.map {
+      case a if a.dataType.isInstanceOf[ArrayType] =>
+        val newArray = ArrayType(a.dataType.asInstanceOf[ArrayType].elementType)
+        val newAttr = AttributeReference(a.name, newArray, a.nullable, a.metadata)()
+        newAttr
+      case other => other
+    }
+    val encoded = ParquetTypesConverter.convertToString(updatedSchama)
     configuration.set(SPARK_ROW_SCHEMA, encoded)
     configuration.set(
       ParquetOutputFormat.WRITER_VERSION,
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
@@ -263,9 +263,9 @@ class HiveContext(sc: SparkContext) extends SQLContext(sc) {
   override protected[sql] lazy val analyzer =
     new Analyzer(catalog, functionRegistry, caseSensitive = false) {
       override val extendedResolutionRules =
+        catalog.PreInsertionCasts ::
         catalog.ParquetConversions ::
         catalog.CreateTables ::
-        catalog.PreInsertionCasts ::
         ExtractPythonUdfs ::
         ResolveUdtfsAlias ::
         sources.PreWriteCheck(catalog) ::
@@ -344,6 +344,7 @@ class HiveContext(sc: SparkContext) extends SQLContext(sc) {
 
     override def strategies: Seq[Strategy] = experimental.extraStrategies ++ Seq(
       DataSourceStrategy,
+      HiveDataSourceStrategy,
       HiveCommandStrategy(self),
       HiveDDLStrategy,
       DDLStrategy,
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveMetastoreCatalog.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveMetastoreCatalog.scala
@@ -424,7 +424,7 @@ private[hive] class HiveMetastoreCatalog(hive: HiveContext) extends Catalog with
       // Collects all `MetastoreRelation`s which should be replaced
       val toBeReplaced = plan.collect {
         // Write path
-        case InsertIntoTable(relation: MetastoreRelation, _, _, _)
+        case InsertIntoHiveTable(relation: MetastoreRelation, _, _, _)
             // Inserting into partitioned table is not supported in Parquet data source (yet).
             if !relation.hiveQlTable.isPartitioned &&
               hive.convertMetastoreParquet &&
@@ -458,6 +458,9 @@ private[hive] class HiveMetastoreCatalog(hive: HiveContext) extends Catalog with
 
           withAlias
         }
+        case InsertIntoHiveTable(r: MetastoreRelation, p, c, o) if relationMap.contains(r) =>
+          val parquetRelation = relationMap(r)
+          InsertIntoHiveTable(parquetRelation, p, c, o) 
         case other => other.transformExpressions {
           case a: Attribute if a.resolved => attributedRewrites.getOrElse(a, a)
         }
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveStrategies.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveStrategies.scala
@@ -34,7 +34,7 @@ import org.apache.spark.sql.execution.{DescribeCommand => RunnableDescribeComman
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.hive.execution._
 import org.apache.spark.sql.parquet.ParquetRelation
-import org.apache.spark.sql.sources.{CreateTableUsingAsSelect, CreateTableUsing}
+import org.apache.spark.sql.sources.{CreateTableUsingAsSelect, CreateTableUsing, LogicalRelation, InsertIntoDataSource, InsertableRelation}
 import org.apache.spark.sql.types.StringType
 
 
@@ -254,4 +254,13 @@ private[hive] trait HiveStrategies {
       case _ => Nil
     }
   }
+
+  object HiveDataSourceStrategy extends Strategy {
+    def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+      case i @ InsertIntoHiveTable(
+        l @ LogicalRelation(t: InsertableRelation), part, query, overwrite) if part.isEmpty =>
+        ExecutedCommand(InsertIntoDataSource(l, query, overwrite)) :: Nil
+      case _ => Nil
+    }
+  }
 }
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/parquet/parquetSuites.scala b/sql/hive/src/test/scala/org/apache/spark/sql/parquet/parquetSuites.scala
@@ -20,6 +20,8 @@ package org.apache.spark.sql.parquet
 
 import java.io.File
 
+import scala.collection.mutable.ArrayBuffer
+
 import org.scalatest.BeforeAndAfterAll
 
 import org.apache.spark.sql.{SQLConf, QueryTest}
@@ -299,6 +301,37 @@ class ParquetDataSourceOnSourceSuite extends ParquetSourceSuiteBase {
     super.afterAll()
     setConf(SQLConf.PARQUET_USE_DATA_SOURCE_API, originalConf.toString)
   }
+
+  test("insert array into parquet hive table using data source api") {
+    val data1="""{ "timestamp": 1422435598, "data_array": [ { "field0": null, "field1": 1, "field2": 2} ] }"""
+    val data2="""{ "timestamp": 1422435599, "data_array": [ { "field0": 0, "field1": null, "field2": 3} ] }"""
+
+    val json = sparkContext.makeRDD(data1 :: data2 :: Nil)
+    val rdd = jsonRDD(json)
+    rdd.registerTempTable("tmp_table")
+
+    val partitionedTableDir = File.createTempFile("persisted_table", "sparksql")
+    partitionedTableDir.delete()
+    partitionedTableDir.mkdir()
+
+    sql(
+      s"""
+        |create external table persisted_table
+        |(
+        |  data_array ARRAY <STRUCT<field0: BIGINT, field1: BIGINT, field2: BIGINT>>,
+        |  timestamp BIGINT
+        |)
+        |STORED AS PARQUET Location '${partitionedTableDir.getCanonicalPath}'
+      """.stripMargin)
+
+    sql("insert into table persisted_table select * from tmp_table").collect
+
+    checkAnswer(
+      sql("select data_array.field0, data_array.field1, data_array.field2 from persisted_table"),
+      Row(ArrayBuffer(null), ArrayBuffer(1), ArrayBuffer(2)) ::
+      Row (ArrayBuffer(0), ArrayBuffer(null), ArrayBuffer(3)) :: Nil
+    )
+  }
 }
 
 class ParquetDataSourceOffSourceSuite extends ParquetSourceSuiteBase {
