Check for empty iterators

ankurdave · ankurdave · commit e27da44e9ce8 · 2014-03-28T15:21:14.000-07:00
diff --git a/graphx/src/main/scala/org/apache/spark/graphx/EdgeRDD.scala b/graphx/src/main/scala/org/apache/spark/graphx/EdgeRDD.scala
@@ -45,7 +45,12 @@ class EdgeRDD[@specialized ED: ClassTag](
     partitionsRDD.partitioner.orElse(Some(Partitioner.defaultPartitioner(partitionsRDD)))
 
   override def compute(part: Partition, context: TaskContext): Iterator[Edge[ED]] = {
-    firstParent[(PartitionID, EdgePartition[ED])].iterator(part, context).next._2.iterator
+    val partIter = firstParent[(PartitionID, EdgePartition[ED])].iterator(part, context)
+    if (partIter.hasNext) {
+      partIter.next._2.iterator
+    } else {
+      Iterator.empty
+    }
   }
 
   override def collect(): Array[Edge[ED]] = this.map(_.copy()).collect()
@@ -69,8 +74,12 @@ class EdgeRDD[@specialized ED: ClassTag](
   private[graphx] def mapEdgePartitions[ED2: ClassTag](
       f: (PartitionID, EdgePartition[ED]) => EdgePartition[ED2]): EdgeRDD[ED2] = {
     new EdgeRDD[ED2](partitionsRDD.mapPartitions({ iter =>
-      val (pid, ep) = iter.next()
-      Iterator(Tuple2(pid, f(pid, ep)))
+      if (iter.hasNext) {
+        val (pid, ep) = iter.next()
+        Iterator(Tuple2(pid, f(pid, ep)))
+      } else {
+        Iterator.empty
+      }
     }, preservesPartitioning = true))
   }
 
@@ -107,9 +116,13 @@ class EdgeRDD[@specialized ED: ClassTag](
     val ed3Tag = classTag[ED3]
     new EdgeRDD[ED3](partitionsRDD.zipPartitions(other.partitionsRDD, true) {
       (thisIter, otherIter) =>
-        val (pid, thisEPart) = thisIter.next()
-        val (_, otherEPart) = otherIter.next()
-        Iterator(Tuple2(pid, thisEPart.innerJoin(otherEPart)(f)(ed2Tag, ed3Tag)))
+        if (thisIter.hasNext && otherIter.hasNext) {
+          val (pid, thisEPart) = thisIter.next()
+          val (_, otherEPart) = otherIter.next()
+          Iterator(Tuple2(pid, thisEPart.innerJoin(otherEPart)(f)(ed2Tag, ed3Tag)))
+        } else {
+          Iterator.empty
+        }
     })
   }
 
diff --git a/graphx/src/main/scala/org/apache/spark/graphx/impl/GraphImpl.scala b/graphx/src/main/scala/org/apache/spark/graphx/impl/GraphImpl.scala
@@ -59,9 +59,13 @@ class GraphImpl[VD: ClassTag, ED: ClassTag] protected (
     val edTag = classTag[ED]
     edges.partitionsRDD.zipPartitions(
       replicatedVertexView.get(true, true), true) { (ePartIter, vPartIter) =>
-      val (pid, ePart) = ePartIter.next()
-      val (_, vPart) = vPartIter.next()
-      new EdgeTripletIterator(vPart.index, vPart.values, ePart)(vdTag, edTag)
+      if (ePartIter.hasNext && vPartIter.hasNext) {
+        val (pid, ePart) = ePartIter.next()
+        val (_, vPart) = vPartIter.next()
+        new EdgeTripletIterator(vPart.index, vPart.values, ePart)(vdTag, edTag)
+      } else {
+        Iterator.empty
+      }
     }
   }
 
@@ -132,22 +136,26 @@ class GraphImpl[VD: ClassTag, ED: ClassTag] protected (
     val newEdgePartitions =
       edges.partitionsRDD.zipPartitions(replicatedVertexView.get(true, true), true) {
         (ePartIter, vTableReplicatedIter) =>
-        val (ePid, edgePartition) = ePartIter.next()
-        val (vPid, vPart) = vTableReplicatedIter.next()
-        assert(!vTableReplicatedIter.hasNext)
-        assert(ePid == vPid)
-        val et = new EdgeTriplet[VD, ED]
-        val inputIterator = edgePartition.iterator.map { e =>
-          et.set(e)
-          et.srcAttr = vPart(e.srcId)
-          et.dstAttr = vPart(e.dstId)
-          et
-        }
-        // Apply the user function to the vertex partition
-        val outputIter = f(ePid, inputIterator)
-        // Consume the iterator to update the edge attributes
-        val newEdgePartition = edgePartition.map(outputIter)
-        Iterator((ePid, newEdgePartition))
+          if (ePartIter.hasNext && vTableReplicatedIter.hasNext) {
+            val (ePid, edgePartition) = ePartIter.next()
+            val (vPid, vPart) = vTableReplicatedIter.next()
+            assert(!vTableReplicatedIter.hasNext)
+            assert(ePid == vPid)
+            val et = new EdgeTriplet[VD, ED]
+            val inputIterator = edgePartition.iterator.map { e =>
+              et.set(e)
+              et.srcAttr = vPart(e.srcId)
+              et.dstAttr = vPart(e.dstId)
+              et
+            }
+            // Apply the user function to the vertex partition
+            val outputIter = f(ePid, inputIterator)
+            // Consume the iterator to update the edge attributes
+            val newEdgePartition = edgePartition.map(outputIter)
+            Iterator((ePid, newEdgePartition))
+          } else {
+            Iterator.empty
+          }
       }
     new GraphImpl(vertices, new EdgeRDD(newEdgePartitions), routingTable, replicatedVertexView)
   }
@@ -217,7 +225,7 @@ class GraphImpl[VD: ClassTag, ED: ClassTag] protected (
 
     // Map and combine.
     val preAgg = edges.partitionsRDD.zipPartitions(vs, true) { (ePartIter, vPartIter) =>
-      if (ePartIter.hasNext) {
+      if (ePartIter.hasNext && vPartIter.hasNext) {
         val (ePid, edgePartition) = ePartIter.next()
         val (vPid, vPart) = vPartIter.next()
         assert(!vPartIter.hasNext)
diff --git a/graphx/src/main/scala/org/apache/spark/graphx/impl/ReplicatedVertexView.scala b/graphx/src/main/scala/org/apache/spark/graphx/impl/ReplicatedVertexView.scala
@@ -103,9 +103,13 @@ class ReplicatedVertexView[VD: ClassTag](
     // Update the view with shippedActives, setting activeness flags in the resulting
     // VertexPartitions
     get(includeSrc, includeDst).zipPartitions(shippedActives) { (viewIter, shippedActivesIter) =>
-      val (pid, vPart) = viewIter.next()
-      val newPart = vPart.replaceActives(shippedActivesIter.flatMap(_._2.iterator))
-      Iterator((pid, newPart))
+      if (viewIter.hasNext) {
+        val (pid, vPart) = viewIter.next()
+        val newPart = vPart.replaceActives(shippedActivesIter.flatMap(_._2.iterator))
+        Iterator((pid, newPart))
+      } else {
+        Iterator.empty
+      }
     }
   }
 
@@ -126,30 +130,40 @@ class ReplicatedVertexView[VD: ClassTag](
         // VertexPartitions
         prevView.get(includeSrc, includeDst).zipPartitions(shippedVerts) {
           (prevViewIter, shippedVertsIter) =>
-            val (pid, prevVPart) = prevViewIter.next()
-            val newVPart = prevVPart.innerJoinKeepLeftDestructive(shippedVertsIter.flatMap(_._2.iterator))
-            Iterator((pid, newVPart))
+            if (prevViewIter.hasNext) {
+              val (pid, prevVPart) = prevViewIter.next()
+              val newVPart =
+                if (destructiveLocal) prevVPart.innerJoinKeepLeftDestructive(shippedVertsIter.flatMap(_._2.iterator))
+                else prevVPart.innerJoinKeepLeft(shippedVertsIter.flatMap(_._2.iterator))
+              Iterator((pid, newVPart))
+            } else {
+              Iterator.empty
+            }
         }.cache().setName("ReplicatedVertexView delta %s %s".format(includeSrc, includeDst))
 
       case None =>
         // Within each edge partition, place the shipped vertex attributes into the correct
         // locations specified in localVertexIdMap
         localVertexIdMap.zipPartitions(shippedVerts) { (mapIter, shippedVertsIter) =>
-          val (pid, vidToIndex) = mapIter.next()
-          assert(!mapIter.hasNext)
-          // Populate the vertex array using the vidToIndex map
-          val vertexArray = vdTag.newArray(vidToIndex.capacity)
-          for ((_, block) <- shippedVertsIter) {
-            for (i <- 0 until block.vids.size) {
-              val vid = block.vids(i)
-              val attr = block.attrs(i)
-              val ind = vidToIndex.getPos(vid)
-              vertexArray(ind) = attr
+          if (mapIter.hasNext) {
+            val (pid, vidToIndex) = mapIter.next()
+            assert(!mapIter.hasNext)
+            // Populate the vertex array using the vidToIndex map
+            val vertexArray = vdTag.newArray(vidToIndex.capacity)
+            for ((_, block) <- shippedVertsIter) {
+              for (i <- 0 until block.vids.size) {
+                val vid = block.vids(i)
+                val attr = block.attrs(i)
+                val ind = vidToIndex.getPos(vid)
+                vertexArray(ind) = attr
+              }
             }
+            val newVPart = new VertexPartition(
+              vidToIndex, vertexArray, vidToIndex.getBitSet)(vdTag)
+            Iterator((pid, newVPart))
+          } else {
+            Iterator.empty
           }
-          val newVPart = new VertexPartition(
-            vidToIndex, vertexArray, vidToIndex.getBitSet)(vdTag)
-          Iterator((pid, newVPart))
         }.cache().setName("ReplicatedVertexView %s %s".format(includeSrc, includeDst))
     }
   }
@@ -159,47 +173,55 @@ private object ReplicatedVertexView {
   protected def buildBuffer[VD: ClassTag](
       pid2vidIter: Iterator[Array[Array[VertexId]]],
       vertexPartIter: Iterator[VertexPartition[VD]]) = {
-    val pid2vid: Array[Array[VertexId]] = pid2vidIter.next()
-    val vertexPart: VertexPartition[VD] = vertexPartIter.next()
-
-    Iterator.tabulate(pid2vid.size) { pid =>
-      val vidsCandidate = pid2vid(pid)
-      val size = vidsCandidate.length
-      val vids = new PrimitiveVector[VertexId](pid2vid(pid).size)
-      val attrs = new PrimitiveVector[VD](pid2vid(pid).size)
-      var i = 0
-      while (i < size) {
-        val vid = vidsCandidate(i)
-        if (vertexPart.isDefined(vid)) {
-          vids += vid
-          attrs += vertexPart(vid)
+    if (pid2vidIter.hasNext && vertexPartIter.hasNext) {
+      val pid2vid: Array[Array[VertexId]] = pid2vidIter.next()
+      val vertexPart: VertexPartition[VD] = vertexPartIter.next()
+
+      Iterator.tabulate(pid2vid.size) { pid =>
+        val vidsCandidate = pid2vid(pid)
+        val size = vidsCandidate.length
+        val vids = new PrimitiveVector[VertexId](pid2vid(pid).size)
+        val attrs = new PrimitiveVector[VD](pid2vid(pid).size)
+        var i = 0
+        while (i < size) {
+          val vid = vidsCandidate(i)
+          if (vertexPart.isDefined(vid)) {
+            vids += vid
+            attrs += vertexPart(vid)
+          }
+          i += 1
         }
-        i += 1
+        (pid, new VertexAttributeBlock(vids.trim().array, attrs.trim().array))
       }
-      (pid, new VertexAttributeBlock(vids.trim().array, attrs.trim().array))
+    } else {
+      Iterator.empty
     }
   }
 
   protected def buildActiveBuffer(
       pid2vidIter: Iterator[Array[Array[VertexId]]],
       activePartIter: Iterator[VertexPartition[_]])
     : Iterator[(Int, Array[VertexId])] = {
-    val pid2vid: Array[Array[VertexId]] = pid2vidIter.next()
-    val activePart: VertexPartition[_] = activePartIter.next()
-
-    Iterator.tabulate(pid2vid.size) { pid =>
-      val vidsCandidate = pid2vid(pid)
-      val size = vidsCandidate.length
-      val actives = new PrimitiveVector[VertexId](vidsCandidate.size)
-      var i = 0
-      while (i < size) {
-        val vid = vidsCandidate(i)
-        if (activePart.isDefined(vid)) {
-          actives += vid
+    if (pid2vidIter.hasNext && activePartIter.hasNext) {
+      val pid2vid: Array[Array[VertexId]] = pid2vidIter.next()
+      val activePart: VertexPartition[_] = activePartIter.next()
+
+      Iterator.tabulate(pid2vid.size) { pid =>
+        val vidsCandidate = pid2vid(pid)
+        val size = vidsCandidate.length
+        val actives = new PrimitiveVector[VertexId](vidsCandidate.size)
+        var i = 0
+        while (i < size) {
+          val vid = vidsCandidate(i)
+          if (activePart.isDefined(vid)) {
+            actives += vid
+          }
+          i += 1
         }
-        i += 1
+        (pid, actives.trim().array)
       }
-      (pid, actives.trim().array)
+    } else {
+      Iterator.empty
     }
   }
 }
