make PartitionPreservingOperationSpec work with Spark 3.0

Signed-off-by: Yitao Li <[email protected]>
twosigma · May 22, 2020 · 0c2529a · 0c2529a
1 parent a51cef0
commit 0c2529a
Showing 1 changed file with 19 additions and 5 deletions.
diff --git a/src/test/scala/org/apache/spark/sql/PartitionPreservingOperationSpec.scala b/src/test/scala/org/apache/spark/sql/PartitionPreservingOperationSpec.scala
@@ -106,17 +106,31 @@ class PartitionPreservingOperationSpec extends FlintSuite with FlintTestData {
     assert(leafExecutedPlan(data).isInstanceOf[ExternalRDDScanExec[_]])
     data.cache()
     data.count()
-    // first node is WholeStageCodegen
-    assert(executedPlan(data).children.head.isInstanceOf[InMemoryTableScanExec])
+    if (data.sparkSession.version < "3.0.0") {
+      // first node is WholeStageCodegen
+      assert(executedPlan(data).children.head.isInstanceOf[InMemoryTableScanExec])
+    } else {
+      val plan = executedPlan(data)
+      assert(plan.children.head.isInstanceOf[ColumnarToRowExec])
+      assert(plan.children.head.children.head.isInstanceOf[InputAdapter])
+      assert(plan.children.head.children.head.children.head.isInstanceOf[InMemoryTableScanExec])
+    }
     assert(leafExecutedPlan(data).isInstanceOf[InMemoryTableScanExec])
     data.unpersist()
 
     val orderedData = data.orderBy("time")
     orderedData.cache()
     orderedData.count()
-    // first node is WholeStageCodegen
-    assert(executedPlan(orderedData).children.head.isInstanceOf[InMemoryTableScanExec])
-    assert(leafExecutedPlan(orderedData).isInstanceOf[InMemoryTableScanExec])
+    if (data.sparkSession.version < "3.0.0") {
+      // first node is WholeStageCodegen
+      assert(executedPlan(orderedData).children.head.isInstanceOf[InMemoryTableScanExec])
+      assert(leafExecutedPlan(orderedData).isInstanceOf[InMemoryTableScanExec])
+    } else {
+      val plan = executedPlan(orderedData)
+      assert(plan.children.head.isInstanceOf[ColumnarToRowExec])
+      assert(plan.children.head.children.head.isInstanceOf[InputAdapter])
+      assert(plan.children.head.children.head.children.head.isInstanceOf[InMemoryTableScanExec])
+    }
     orderedData.unpersist()
   }