I really need to stop calling size on iterators, it is the path of sa…

…dness.
pdeyhim · Apr 8, 2014 · 71e8b9f · 71e8b9f
1 parent b1ae51a
commit 71e8b9f
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/bagel/src/main/scala/org/apache/spark/bagel/Bagel.scala b/bagel/src/main/scala/org/apache/spark/bagel/Bagel.scala
@@ -228,12 +228,12 @@ object Bagel extends Logging {
     var numActiveVerts = sc.accumulator(0)
     val processed = grouped.flatMapValues {
       case (_, vs) if !vs.hasNext => None
-      case (c, vs) =>
+      case (c, vs) => {
         val (newVert, newMsgs) =
           compute(vs.next,
-            c.size match {
-              case 1 => Some(c.next)
-              case _ => None
+            c.hasNext match {
+              case true => Some(c.next)
+              case false => None
             }
           )
 
@@ -243,6 +243,7 @@ object Bagel extends Logging {
         }
 
         Some((newVert, newMsgs))
+      }
     }.persist(storageLevel)
 
     // Force evaluation of processed RDD for accurate performance measurements