restore the SparkEnv after SparkContext.stop()

Author: chenghao-intel

sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeRowSerializerSuite.scala

@@ -22,6 +22,7 @@ import java.io.{File, DataOutputStream, ByteArrayInputStream, ByteArrayOutputStr
 import org.apache.spark.executor.ShuffleWriteMetrics
 import org.apache.spark.storage.ShuffleBlockId
 import org.apache.spark.util.collection.ExternalSorter
+import org.apache.spark.util.Utils
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow}
 import org.apache.spark.sql.catalyst.expressions.{UnsafeProjection, UnsafeRow}
@@ -43,9 +44,15 @@ class ClosableByteArrayInputStream(buf: Array[Byte]) extends ByteArrayInputStrea
 class UnsafeRowSerializerSuite extends SparkFunSuite {
 
   private def toUnsafeRow(row: Row, schema: Array[DataType]): UnsafeRow = {
-    val internalRow = CatalystTypeConverters.convertToCatalyst(row).asInstanceOf[InternalRow]
+    val converter = unsafeRowConverter(schema)
+    converter(row)
+  }
+
+  private def unsafeRowConverter(schema: Array[DataType]): Row => UnsafeRow = {
     val converter = UnsafeProjection.create(schema)
-    converter.apply(internalRow)
+    (row: Row) => {
+      converter(CatalystTypeConverters.convertToCatalyst(row).asInstanceOf[InternalRow])
+    }
   }
 
   test("toUnsafeRow() test helper method") {
@@ -92,37 +99,45 @@ class UnsafeRowSerializerSuite extends SparkFunSuite {
   }
 
   test("SPARK-10466: external sorter spilling with unsafe row serializer") {
-    val conf = new SparkConf()
-      .set("spark.shuffle.spill.initialMemoryThreshold", "1024")
-      .set("spark.shuffle.sort.bypassMergeThreshold", "0")
-      .set("spark.shuffle.memoryFraction", "0.0001")
     var sc: SparkContext = null
     var outputFile: File = null
-    try {
-      sc = new SparkContext("local", "test", conf)
-      outputFile = File.createTempFile("test-unsafe-row-serializer-spill", "")
-      val data = (1 to 1000).iterator.map { i =>
-        (i, toUnsafeRow(Row(i), Array(IntegerType)))
-      }
-      val sorter = new ExternalSorter[Int, UnsafeRow, UnsafeRow](
-        partitioner = Some(new HashPartitioner(10)),
-        serializer = Some(new UnsafeRowSerializer(numFields = 1)))
+    val oldEnv = SparkEnv.get // save the old SparkEnv, as it will be overwritten
+    Utils.tryWithSafeFinally {
+      val conf = new SparkConf()
+        .set("spark.shuffle.spill.initialMemoryThreshold", "1024")
+        .set("spark.shuffle.sort.bypassMergeThreshold", "0")
+        .set("spark.shuffle.memoryFraction", "0.0001")
 
-      // Ensure we spilled something and have to merge them later
-      assert(sorter.numSpills === 0)
-      sorter.insertAll(data)
-      assert(sorter.numSpills > 0)
+        sc = new SparkContext("local", "test", conf)
+        outputFile = File.createTempFile("test-unsafe-row-serializer-spill", "")
+        // prepare data
+        val converter = unsafeRowConverter(Array(IntegerType))
+        val data = (1 to 1000).iterator.map { i =>
+          (i, converter(Row(i)))
+        }
+        val sorter = new ExternalSorter[Int, UnsafeRow, UnsafeRow](
+          partitioner = Some(new HashPartitioner(10)),
+          serializer = Some(new UnsafeRowSerializer(numFields = 1)))
 
-      // Merging spilled files should not throw assertion error
-      val taskContext = new TaskContextImpl(0, 0, 0, 0, null, null, InternalAccumulator.create(sc))
-      taskContext.taskMetrics.shuffleWriteMetrics = Some(new ShuffleWriteMetrics)
-      sorter.writePartitionedFile(ShuffleBlockId(0, 0, 0), taskContext, outputFile)
+        // Ensure we spilled something and have to merge them later
+        assert(sorter.numSpills === 0)
+        sorter.insertAll(data)
+        assert(sorter.numSpills > 0)
 
-    } finally {
+        // Merging spilled files should not throw assertion error
+        val taskContext =
+          new TaskContextImpl(0, 0, 0, 0, null, null, InternalAccumulator.create(sc))
+        taskContext.taskMetrics.shuffleWriteMetrics = Some(new ShuffleWriteMetrics)
+        sorter.writePartitionedFile(ShuffleBlockId(0, 0, 0), taskContext, outputFile)
+    } {
       // Clean up
       if (sc != null) {
         sc.stop()
       }
+
+      // restore the spark env
+      SparkEnv.set(oldEnv)
+
       if (outputFile != null) {
         outputFile.delete()
       }