Merge pull request #17 from amplab/product2

product 2 change

graph/src/main/scala/org/apache/spark/graph/GraphKryoRegistrator.scala

@@ -2,23 +2,19 @@ package org.apache.spark.graph
 
 import com.esotericsoftware.kryo.Kryo
 
+import org.apache.spark.graph.impl.MessageToPartition
 import org.apache.spark.serializer.KryoRegistrator
 
 
 class GraphKryoRegistrator extends KryoRegistrator {
 
   def registerClasses(kryo: Kryo) {
-    //kryo.register(classOf[(Int, Float, Float)])
-    registerClass[Int, Int, Int](kryo)
+    kryo.register(classOf[Vertex[Object]])
+    kryo.register(classOf[Edge[Object]])
+    kryo.register(classOf[MutableTuple2[Object, Object]])
+    kryo.register(classOf[MessageToPartition[Object]])
 
     // This avoids a large number of hash table lookups.
     kryo.setReferences(false)
   }
-
-  private def registerClass[VD: Manifest, ED: Manifest, VD2: Manifest](kryo: Kryo) {
-    kryo.register(classOf[Vertex[VD]])
-    kryo.register(classOf[Edge[ED]])
-    kryo.register(classOf[MutableTuple2[VD, VD2]])
-    kryo.register(classOf[(Vid, VD2)])
-  }
 }

graph/src/main/scala/org/apache/spark/graph/impl/GraphImpl.scala

@@ -11,9 +11,7 @@ import org.apache.spark.rdd.RDD
 
 import org.apache.spark.graph._
 import org.apache.spark.graph.impl.GraphImpl._
-
-
-
+import org.apache.spark.graph.impl.MessageToPartitionRDDFunctions._
 
 
 /**
@@ -309,11 +307,13 @@ class GraphImpl[VD: ClassManifest, ED: ClassManifest] protected (
     // Join vid2pid and vTable, generate a shuffle dependency on the joined result, and get
     // the shuffle id so we can use it on the slave.
     vTable
-      .flatMap { case (vid, (vdata, pids)) => pids.iterator.map { pid => (pid, (vid, vdata)) } }
+      .flatMap { case (vid, (vdata, pids)) =>
+        pids.iterator.map { pid => MessageToPartition(pid, (vid, vdata)) }
+      }
       .partitionBy(edgePartitioner)
-      .mapPartitions(
-        { part => part.map { case(pid, (vid, vdata)) => (vid, vdata) } },
-        preservesPartitioning = true)
+      .mapPartitions({ part =>
+        part.map { message => (message.data._1, message.data._2) }
+      }, preservesPartitioning = true)
   }
 }
 
@@ -400,14 +400,16 @@ object GraphImpl {
         val part: Pid = edgePartitionFunction2D(e.src, e.dst, numPartitions, ceilSqrt)
 
         // Should we be using 3-tuple or an optimized class
-        (part, (e.src, e.dst, e.data))
+        MessageToPartition(part, (e.src, e.dst, e.data))
         //  (math.abs(e.src) % numPartitions, (e.src, e.dst, e.data))
-       
       }
       .partitionBy(new HashPartitioner(numPartitions))
       .mapPartitionsWithIndex({ (pid, iter) =>
         val edgePartition = new EdgePartition[ED]
-        iter.foreach { case (_, (src, dst, data)) => edgePartition.add(src, dst, data) }
+        iter.foreach { message =>
+          val data = message.data
+          edgePartition.add(data._1, data._2, data._3)
+        }
         edgePartition.trim()
         Iterator((pid, edgePartition))
       }, preservesPartitioning = true)

graph/src/main/scala/org/apache/spark/graph/impl/MessageToPartition.scala

@@ -0,0 +1,49 @@
+package org.apache.spark.graph.impl
+
+import org.apache.spark.Partitioner
+import org.apache.spark.graph.Pid
+import org.apache.spark.rdd.{ShuffledRDD, RDD}
+
+
+/**
+ * A message used to send a specific value to a partition.
+ * @param partition index of the target partition.
+ * @param data value to send
+ */
+class MessageToPartition[@specialized(Int, Long, Double, Char, Boolean/*, AnyRef*/) T](
+    @transient var partition: Pid,
+    var data: T)
+  extends Product2[Pid, T] {
+
+  override def _1 = partition
+
+  override def _2 = data
+
+  override def canEqual(that: Any): Boolean = that.isInstanceOf[MessageToPartition[_]]
+}
+
+/**
+ * Companion object for MessageToPartition.
+ */
+object MessageToPartition {
+  def apply[T](partition: Pid, value: T) = new MessageToPartition(partition, value)
+}
+
+
+class MessageToPartitionRDDFunctions[T: ClassManifest](self: RDD[MessageToPartition[T]]) {
+
+  /**
+   * Return a copy of the RDD partitioned using the specified partitioner.
+   */
+  def partitionBy(partitioner: Partitioner): RDD[MessageToPartition[T]] = {
+    new ShuffledRDD[Pid, T, MessageToPartition[T]](self, partitioner)
+  }
+
+}
+
+
+object MessageToPartitionRDDFunctions {
+  implicit def rdd2PartitionRDDFunctions[T: ClassManifest](rdd: RDD[MessageToPartition[T]]) = {
+    new MessageToPartitionRDDFunctions(rdd)
+  }
+}