Refactor and add aggregator support

Refactored out the agg() and comp() methods from Pregel.run.

Defined an implicit conversion to allow applications that don't use
aggregators to avoid including a null argument for the result of the
aggregator in the compute function.

1 files changed, 6 insertions, 4 deletions

diff --git a/bagel/src/test/scala/bagel/BagelSuite.scala b/bagel/src/test/scala/bagel/BagelSuite.scala
index 29f5f0c358..53a93a6b80 100644
--- a/bagel/src/test/scala/bagel/BagelSuite.scala
+++ b/bagel/src/test/scala/bagel/BagelSuite.scala
@@ -10,6 +10,8 @@ import scala.collection.mutable.ArrayBuffer
 
 import spark._
 
+import bagel.Pregel._
+
 @serializable class TestVertex(val id: String, val active: Boolean, val age: Int) extends Vertex
 @serializable class TestMessage(val targetId: String) extends Message
 
@@ -20,10 +22,10 @@ class BagelSuite extends FunSuite with Assertions {
     val msgs = sc.parallelize(Array[(String, TestMessage)]())
     val numSupersteps = 5
     val result =
-      Pregel.run(sc, verts, msgs, new DefaultCombiner[TestMessage], 1) {
+      Pregel.run(sc, verts, msgs)()(addAggregatorArg {
       (self: TestVertex, msgs: Option[ArrayBuffer[TestMessage]], superstep: Int) =>
         (new TestVertex(self.id, superstep < numSupersteps - 1, self.age + 1), Array[TestMessage]())
-    }
+    })
     for (vert <- result.collect)
       assert(vert.age === numSupersteps)
   }
@@ -34,7 +36,7 @@ class BagelSuite extends FunSuite with Assertions {
     val msgs = sc.parallelize(Array("a" -> new TestMessage("a")))
     val numSupersteps = 5
     val result =
-      Pregel.run(sc, verts, msgs, new DefaultCombiner[TestMessage], 1) {
+      Pregel.run(sc, verts, msgs)()(addAggregatorArg {
       (self: TestVertex, msgs: Option[ArrayBuffer[TestMessage]], superstep: Int) =>
         val msgsOut =
           msgs match {
@@ -44,7 +46,7 @@ class BagelSuite extends FunSuite with Assertions {
               new ArrayBuffer[TestMessage]()
           }
         (new TestVertex(self.id, self.active, self.age + 1), msgsOut)
-    }
+    })
     for (vert <- result.collect)
       assert(vert.age === numSupersteps)
   }