Fix a couple tests and move getAutoReset to KryoSerializerInstance

Author: sryza

core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala

@@ -124,17 +124,6 @@ class KryoSerializer(conf: SparkConf)
   override def newInstance(): SerializerInstance = {
     new KryoSerializerInstance(this)
   }
-
-  /**
-   * Returns true if auto-reset is on. The only reason this would be false is if the user-supplied
-   * register explicitly turns auto-reset off.
-   */
-  def getAutoReset(): Boolean = {
-    val kryo = newKryo()
-    val field = classOf[Kryo].getDeclaredField("autoReset")
-    field.setAccessible(true)
-    field.get(kryo).asInstanceOf[Boolean]
-  }
 }
 
 private[spark]
@@ -210,6 +199,16 @@ private[spark] class KryoSerializerInstance(ks: KryoSerializer) extends Serializ
   override def deserializeStream(s: InputStream): DeserializationStream = {
     new KryoDeserializationStream(kryo, s)
   }
+
+  /**
+   * Returns true if auto-reset is on. The only reason this would be false is if the user-supplied
+   * registrator explicitly turns auto-reset off.
+   */
+  def getAutoReset(): Boolean = {
+    val field = classOf[Kryo].getDeclaredField("autoReset")
+    field.setAccessible(true)
+    field.get(kryo).asInstanceOf[Boolean]
+  }
 }
 
 /**

core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala

@@ -26,7 +26,7 @@ import scala.collection.mutable
 import com.google.common.io.ByteStreams
 
 import org.apache.spark._
-import org.apache.spark.serializer.{DeserializationStream, KryoSerializer, Serializer}
+import org.apache.spark.serializer._
 import org.apache.spark.executor.ShuffleWriteMetrics
 import org.apache.spark.storage.{BlockObjectWriter, BlockId}
 
@@ -129,14 +129,15 @@ private[spark] class ExternalSorter[K, V, C](
   private val kvChunkSize = conf.getInt("spark.shuffle.sort.kvChunkSize", 1 << 22) // 4 MB
   private val useSerializedPairBuffer =
     !ordering.isDefined && conf.getBoolean("spark.shuffle.sort.serializeMapOutputs", true) &&
-    ser.isInstanceOf[KryoSerializer] && ser.asInstanceOf[KryoSerializer].getAutoReset
+    ser.isInstanceOf[KryoSerializer] &&
+    serInstance.asInstanceOf[KryoSerializerInstance].getAutoReset
 
   // Data structures to store in-memory objects before we spill. Depending on whether we have an
   // Aggregator set, we either put objects into an AppendOnlyMap where we combine them, or we
   // store them in an array buffer.
   private var map = new PartitionedAppendOnlyMap[K, C]
   private var buffer = if (useSerializedPairBuffer) {
-    new PartitionedSerializedPairBuffer[K, C](metaInitialRecords, kvChunkSize)
+    new PartitionedSerializedPairBuffer[K, C](metaInitialRecords, kvChunkSize, serInstance)
   } else {
     new PartitionedPairBuffer[K, C]
   }
@@ -237,7 +238,7 @@ private[spark] class ExternalSorter[K, V, C](
     } else {
       if (maybeSpill(buffer, buffer.estimateSize())) {
         buffer = if (useSerializedPairBuffer) {
-          new PartitionedSerializedPairBuffer[K, C](metaInitialRecords, kvChunkSize)
+          new PartitionedSerializedPairBuffer[K, C](metaInitialRecords, kvChunkSize, serInstance)
         } else {
           new PartitionedPairBuffer[K, C]
         }

core/src/main/scala/org/apache/spark/util/collection/PartitionedSerializedPairBuffer.scala

@@ -48,7 +48,7 @@ import org.apache.spark.util.collection.PartitionedSerializedPairBuffer._
 private[spark] class PartitionedSerializedPairBuffer[K, V](
     metaInitialRecords: Int,
     kvBlockSize: Int,
-    serializerInstance: SerializerInstance = SparkEnv.get.serializer.newInstance)
+    serializerInstance: SerializerInstance)
   extends WritablePartitionedPairCollection[K, V] with SizeTracker {
 
   if (serializerInstance.isInstanceOf[JavaSerializerInstance]) {

core/src/test/scala/org/apache/spark/serializer/KryoSerializerSuite.scala

@@ -282,11 +282,11 @@ class KryoSerializerSuite extends FunSuite with SharedSparkContext {
   }
 
   test("getAutoReset") {
-    val ser = new KryoSerializer(new SparkConf)
+    val ser = new KryoSerializer(new SparkConf).newInstance().asInstanceOf[KryoSerializerInstance]
     assert(ser.getAutoReset)
     val conf = new SparkConf().set("spark.kryo.registrator",
       classOf[RegistratorWithoutAutoReset].getName)
-    val ser2 = new KryoSerializer(conf)
+    val ser2 = new KryoSerializer(conf).newInstance().asInstanceOf[KryoSerializerInstance]
     assert(!ser2.getAutoReset)
   }
 }

core/src/test/scala/org/apache/spark/util/collection/ChainedBufferSuite.scala

@@ -52,8 +52,12 @@ class ChainedBufferSuite extends FunSuite {
   }
 
   test("write and read at middle") {
-    // write from start of source array
     val buffer = new ChainedBuffer(8)
+
+    // fill to a middle point
+    verifyWriteAndRead(buffer, 0, 0, 0, 3)
+
+    // write from start of source array
     verifyWriteAndRead(buffer, 3, 0, 0, 4)
     buffer.capacity should be (8)
 
@@ -79,8 +83,12 @@ class ChainedBufferSuite extends FunSuite {
   }
 
   test("write and read at later buffer") {
-    // write from start of source array
     val buffer = new ChainedBuffer(8)
+
+    // fill to a middle point
+    verifyWriteAndRead(buffer, 0, 0, 0, 11)
+
+    // write from start of source array
     verifyWriteAndRead(buffer, 11, 0, 0, 4)
     buffer.capacity should be (16)
 

core/src/test/scala/org/apache/spark/util/collection/PartitionedSerializedPairBufferSuite.scala

@@ -52,7 +52,7 @@ class PartitionedSerializedPairBufferSuite extends FunSuite {
     val buffer = new PartitionedSerializedPairBuffer[Int, SomeStruct](4, 32, serializerInstance)
     val struct = SomeStruct("something", 5)
     buffer.insert(4, 10, struct)
-    val elements = buffer.partitionedDestructiveSortedIterator(null).toArray
+    val elements = buffer.partitionedDestructiveSortedIterator(None).toArray
     elements.size should be (1)
     elements.head should be (((4, 10), struct))
   }
@@ -67,7 +67,7 @@ class PartitionedSerializedPairBufferSuite extends FunSuite {
     val struct3 = SomeStruct("something3", 10)
     buffer.insert(5, 3, struct3)
 
-    val elements = buffer.partitionedDestructiveSortedIterator(null).toArray
+    val elements = buffer.partitionedDestructiveSortedIterator(None).toArray
     elements.size should be (3)
     elements(0) should be (((4, 2), struct2))
     elements(1) should be (((5, 3), struct3))
@@ -79,7 +79,7 @@ class PartitionedSerializedPairBufferSuite extends FunSuite {
     val buffer = new PartitionedSerializedPairBuffer[Int, SomeStruct](4, 32, serializerInstance)
     val struct = SomeStruct("something", 5)
     buffer.insert(4, 10, struct)
-    val it = buffer.destructiveSortedWritablePartitionedIterator(null)
+    val it = buffer.destructiveSortedWritablePartitionedIterator(None)
     val writer = new SimpleBlockObjectWriter
     assert(it.hasNext)
     it.nextPartition should be (4)
@@ -101,7 +101,7 @@ class PartitionedSerializedPairBufferSuite extends FunSuite {
     val struct3 = SomeStruct("something3", 10)
     buffer.insert(5, 3, struct3)
 
-    val it = buffer.destructiveSortedWritablePartitionedIterator(null)
+    val it = buffer.destructiveSortedWritablePartitionedIterator(None)
     val writer = new SimpleBlockObjectWriter
     assert(it.hasNext)
     it.nextPartition should be (4)
@@ -142,7 +142,7 @@ class SimpleBlockObjectWriter extends BlockObjectWriter(null) {
   override def isOpen: Boolean = true
   override def commitAndClose(): Unit = { }
   override def revertPartialWritesAndClose(): Unit = { }
-  override def fileSegment(): FileSegment = { null }
+  override def fileSegment(): FileSegment = null
   override def write(key: Any, value: Any): Unit = { }
   override def recordWritten(): Unit = { }
   override def write(b: Int): Unit = { }