Make to local iterator shorter

witgo · Apr 5, 2014 · 8ec8f24 · 8ec8f24
1 parent 34aa300
commit 8ec8f24
Showing 1 changed file with 1 addition and 7 deletions.
diff --git a/core/src/main/scala/org/apache/spark/rdd/RDD.scala b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
@@ -672,13 +672,7 @@ abstract class RDD[T: ClassTag](
     def collectPartition(p: Int): Array[T] = {
       sc.runJob(this, (iter: Iterator[T]) => iter.toArray, Seq(p), allowLocal = false).head
     }
-    var buffer = Stream.empty[T]
-    for (p <- 0 until this.partitions.length) {
-      buffer = buffer append {
-        collectPartition(p).toStream
-      }
-    }
-    buffer.iterator
+    (0 until partitions.length).iterator.flatMap(i => collectPartition(i))
   }
 
   /**