Add ore tests

apache · Jun 10, 2020 · 869330d · 869330d
1 parent caf0ce4
commit 869330d
Showing 1 changed file with 29 additions and 0 deletions.
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala
@@ -69,6 +69,7 @@ class WholeStageCodegenSuite extends QueryTest with SharedSparkSession
       Array(Row("James", "Java", Map("hair" -> "black", "eye" -> "brown")),
         Row("James", "Scala", Map("hair" -> "black", "eye" -> "brown"))))
 
+
     // Map - explode
     expDF = df.select($"name", $"knownLanguages", explode($"properties"))
     plan = expDF.queryExecution.executedPlan
@@ -105,6 +106,34 @@ class WholeStageCodegenSuite extends QueryTest with SharedSparkSession
     results = expDF.collect()
     assert(results ===
       Array(Row("James", 0, "hair", "black"), Row("James", 1, "eye", "brown")))
+
+
+    // Array - explode , selecting all columns
+    expDF = df.select($"*", explode($"knownLanguages"))
+    plan = expDF.queryExecution.executedPlan
+    assert(plan.find {
+      case stage: WholeStageCodegenExec =>
+        stage.find(_.isInstanceOf[GenerateExec]).isDefined
+      case _ => false
+    }.isDefined)
+    results = expDF.collect()
+    assert(results ===
+      Array(Row("James", Seq("Java", "Scala"), Map("hair" -> "black", "eye" -> "brown"), "Java"),
+        Row("James", Seq("Java", "Scala"), Map("hair" -> "black", "eye" -> "brown"), "Scala")))
+
+    // Map - explode, selecting all columns
+    expDF = df.select($"*", explode($"properties"))
+    plan = expDF.queryExecution.executedPlan
+    assert(plan.find {
+      case stage: WholeStageCodegenExec =>
+        stage.find(_.isInstanceOf[GenerateExec]).isDefined
+      case _ => false
+    }.isDefined)
+    results = expDF.collect()
+    assert(results ===
+      Array(
+        Row("James", List("Java", "Scala"), Map("hair" -> "black", "eye" -> "brown"), "hair", "black"),
+        Row("James", List("Java", "Scala"), Map("hair" -> "black", "eye" -> "brown"), "eye", "brown")))
   }
 
   test("Aggregate with grouping keys should be included in WholeStageCodegen") {