Adds a vector serializer

twitter · May 1, 2012 · 9124b74 · 9124b74
1 parent 7aa8979
commit 9124b74
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 0 deletions.
diff --git a/src/main/scala/com/twitter/scalding/KryoHadoopSerialization.scala b/src/main/scala/com/twitter/scalding/KryoHadoopSerialization.scala
@@ -54,6 +54,7 @@ class KryoHadoopSerialization extends KryoSerialization {
   override def decorateKryo(newK : Kryo) : Kryo = {
 
     newK.addDefaultSerializer(classOf[List[Any]], new ListSerializer(List[AnyRef]()))
+    newK.addDefaultSerializer(classOf[Vector[Any]], new VectorSerializer[Any])
     newK.register(classOf[RichDate], new RichDateSerializer())
     newK.register(classOf[DateRange], new DateRangeSerializer())
     // Add some maps
@@ -146,6 +147,34 @@ class ListSerializer[T <: List[_]](emptyList : List[_]) extends KSerializer[T] {
   }
 }
 
+class VectorSerializer[T] extends KSerializer[Vector[T]] {
+  def write(kser: Kryo, out: Output, obj: Vector[T]) {
+    //Write the size:
+    out.writeInt(obj.size, true)
+    obj.foreach { (t : Any) => kser.writeClassAndObject(out, t) }
+  }
+
+  override def create(kser: Kryo, in: Input, cls: Class[Vector[T]]) : Vector[T] = {
+    val size = in.readInt(true);
+
+    //Produce the reversed list:
+    if (size == 0) {
+      /*
+       * this is only here at compile time.  The type T is erased, but the
+       * compiler verifies that we are intending to return a type T here.
+       */
+      Vector.empty[T]
+    }
+    else {
+      (0 until size).foldLeft(Vector.empty[T]) { (vec, i) =>
+        val iT = kser.readClassAndObject(in).asInstanceOf[T]
+        vec :+ iT
+      }
+    }
+  }
+}
+
+
 class MapSerializer[T <: Map[_,_]](emptyMap : Map[_,_]) extends KSerializer[T] {
   def write(kser: Kryo, out: Output, obj: T) {
     //Write the size:

diff --git a/src/test/scala/com/twitter/scalding/KryoTest.scala b/src/test/scala/com/twitter/scalding/KryoTest.scala
@@ -66,6 +66,7 @@ class KryoTest extends Specification {
                       TestCaseClassForSerialization("case classes are: ", 10),
                       TestValMap(Map("you" -> 1.0, "every" -> 2.0, "body" -> 3.0, "a" -> 1.0,
                         "b" -> 2.0, "c" -> 3.0, "d" -> 4.0)),
+                      Vector(1,2,3,4,5),
                       TestValMap(null),
                       Some("junk"))
         .asInstanceOf[List[AnyRef]]