Increase coverage of inserts into a JSONRelation

Author: Nathan Howell

sql/core/src/test/scala/org/apache/spark/sql/sources/InsertSuite.scala

@@ -21,7 +21,7 @@ import java.io.File
 
 import org.scalatest.BeforeAndAfterAll
 
-import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.{SaveMode, AnalysisException, Row}
 import org.apache.spark.util.Utils
 
 class InsertSuite extends DataSourceTest with BeforeAndAfterAll {
@@ -100,23 +100,48 @@ class InsertSuite extends DataSourceTest with BeforeAndAfterAll {
   test("INSERT OVERWRITE a JSONRelation multiple times") {
     sql(
       s"""
-        |INSERT OVERWRITE TABLE jsonTable SELECT a, b FROM jt
-      """.stripMargin)
+         |INSERT OVERWRITE TABLE jsonTable SELECT a, b FROM jt
+    """.stripMargin)
+    checkAnswer(
+      sql("SELECT a, b FROM jsonTable"),
+      (1 to 10).map(i => Row(i, s"str$i"))
+    )
 
+    // Writing the table to less part files.
+    val rdd1 = sparkContext.parallelize((1 to 10).map(i => s"""{"a":$i, "b":"str${i}"}"""), 5)
+    jsonRDD(rdd1).registerTempTable("jt1")
     sql(
       s"""
-        |INSERT OVERWRITE TABLE jsonTable SELECT a, b FROM jt
-      """.stripMargin)
+         |INSERT OVERWRITE TABLE jsonTable SELECT a, b FROM jt1
+    """.stripMargin)
+    checkAnswer(
+      sql("SELECT a, b FROM jsonTable"),
+      (1 to 10).map(i => Row(i, s"str$i"))
+    )
 
+    // Writing the table to more part files.
+    val rdd2 = sparkContext.parallelize((1 to 10).map(i => s"""{"a":$i, "b":"str${i}"}"""), 10)
+    jsonRDD(rdd2).registerTempTable("jt2")
     sql(
       s"""
-        |INSERT OVERWRITE TABLE jsonTable SELECT a, b FROM jt
-      """.stripMargin)
-
+         |INSERT OVERWRITE TABLE jsonTable SELECT a, b FROM jt2
+    """.stripMargin)
     checkAnswer(
       sql("SELECT a, b FROM jsonTable"),
       (1 to 10).map(i => Row(i, s"str$i"))
     )
+
+    sql(
+      s"""
+         |INSERT OVERWRITE TABLE jsonTable SELECT a * 10, b FROM jt1
+    """.stripMargin)
+    checkAnswer(
+      sql("SELECT a, b FROM jsonTable"),
+      (1 to 10).map(i => Row(i * 10, s"str$i"))
+    )
+
+    dropTempTable("jt1")
+    dropTempTable("jt2")
   }
 
   test("INSERT INTO not supported for JSONRelation for now") {
@@ -128,6 +153,20 @@ class InsertSuite extends DataSourceTest with BeforeAndAfterAll {
     }
   }
 
+  test("save directly to the path of a JSON table") {
+    table("jt").selectExpr("a * 5 as a", "b").save(path.toString, "json", SaveMode.Overwrite)
+    checkAnswer(
+      sql("SELECT a, b FROM jsonTable"),
+      (1 to 10).map(i => Row(i * 5, s"str$i"))
+    )
+
+    table("jt").save(path.toString, "json", SaveMode.Overwrite)
+    checkAnswer(
+      sql("SELECT a, b FROM jsonTable"),
+      (1 to 10).map(i => Row(i, s"str$i"))
+    )
+  }
+
   test("it is not allowed to write to a table while querying it.") {
     val message = intercept[AnalysisException] {
       sql(