f in foreach could be generator or not.

apache · Aug 24, 2014 · 1789cd4 · 1789cd4
1 parent 2871b80
commit 1789cd4
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/python/pyspark/rdd.py b/python/pyspark/rdd.py
@@ -655,12 +655,14 @@ def foreachPartition(self, f):
         >>> def f(iterator):
         ...      for x in iterator:
         ...           print x
-        ...      yield None
         >>> sc.parallelize([1, 2, 3, 4, 5]).foreachPartition(f)
         """
         def func(it):
-            f(it)
-            return iter([])
+            r = f(it)
+            try:
+                return iter(r)
+            except TypeError:
+                return iter([])
         self.mapPartitions(func).count()  # Force evaluation
 
     def collect(self):