Optimization for count()

2011-05-13 10:41:34 -07:00 · 2011-05-13 10:41:34 -07:00 · 16c886a581
parent 4b1f0f1ce4
commit 16c886a581
1 changed files with 8 additions and 1 deletions
--- a/core/src/main/scala/spark/RDD.scala
+++ b/core/src/main/scala/spark/RDD.scala
@ -107,7 +107,14 @@ abstract class RDD[T: ClassManifest](@transient sc: SparkContext) {
  }
  
  def count(): Long = {
-    sc.runJob(this, (iter: Iterator[T]) => iter.size.toLong).sum
+    sc.runJob(this, (iter: Iterator[T]) => {
+      var result = 0L
+      while (iter.hasNext) {
+        result += 1L
+        iter.next
+      }
+      result
+    }).sum
  }

  def toArray(): Array[T] = collect()