Register OpenHashMapBasedStateMap for Kryo

Author: zsxwing

streaming/src/main/scala/org/apache/spark/streaming/util/StateMap.scala

@@ -21,6 +21,9 @@ import java.io.{ObjectInputStream, ObjectOutputStream}
 
 import scala.reflect.ClassTag
 
+import com.esotericsoftware.kryo.DefaultSerializer
+import com.esotericsoftware.kryo.serializers.{JavaSerializer => KryoJavaSerializer}
+
 import org.apache.spark.SparkConf
 import org.apache.spark.streaming.util.OpenHashMapBasedStateMap._
 import org.apache.spark.util.collection.OpenHashMap
@@ -77,6 +80,7 @@ private[streaming] class EmptyStateMap[K: ClassTag, S: ClassTag] extends StateMa
 }
 
 /** Implementation of StateMap based on Spark's [[org.apache.spark.util.collection.OpenHashMap]] */
+@DefaultSerializer(classOf[KryoJavaSerializer])
 private[streaming] class OpenHashMapBasedStateMap[K: ClassTag, S: ClassTag](
     @transient @volatile var parentStateMap: StateMap[K, S],
     initialCapacity: Int = DEFAULT_INITIAL_CAPACITY,

streaming/src/test/scala/org/apache/spark/streaming/StateMapSuite.scala

@@ -18,14 +18,17 @@
 package org.apache.spark.streaming
 
 import scala.collection.{immutable, mutable, Map}
+import scala.reflect.ClassTag
 import scala.util.Random
 
-import org.apache.spark.SparkFunSuite
+import org.apache.spark.{SparkConf, SparkFunSuite}
+import org.apache.spark.serializer._
 import org.apache.spark.streaming.util.{EmptyStateMap, OpenHashMapBasedStateMap, StateMap}
-import org.apache.spark.util.Utils
 
 class StateMapSuite extends SparkFunSuite {
 
+  private val conf = new SparkConf()
+
   test("EmptyStateMap") {
     val map = new EmptyStateMap[Int, Int]
     intercept[scala.NotImplementedError] {
@@ -267,12 +270,17 @@ class StateMapSuite extends SparkFunSuite {
     assertMap(stateMap, refMap.toMap, time, "Final state map does not match reference map")
   }
 
-  private def testSerialization[MapType <: StateMap[Int, Int]](
+  private def testSerialization[MapType <: StateMap[Int, Int] : ClassTag](
     map: MapType, msg: String): MapType = {
-    val deserMap = Utils.deserialize[MapType](
-      Utils.serialize(map), Thread.currentThread().getContextClassLoader)
-    assertMap(deserMap, map, 1, msg)
-    deserMap
+    val deserMaps = Array(new JavaSerializer(conf), new KryoSerializer(conf)).map {
+      (serializer: Serializer) =>
+        val serializerInstance = serializer.newInstance()
+        val deserMap = serializerInstance.deserialize(
+          serializerInstance.serialize(map), Thread.currentThread().getContextClassLoader)
+        assertMap(deserMap, map, 1, msg)
+        deserMap
+    }
+    deserMaps.head
   }
 
   // Assert whether all the data and operations on a state map matches that of a reference state map