add wrapper for foreachRDD()

markhamstra · Oct 10, 2014 · 02d0575 · 02d0575
1 parent bebeb4a
commit 02d0575
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/python/pyspark/streaming/dstream.py b/python/pyspark/streaming/dstream.py
@@ -150,6 +150,9 @@ def foreachRDD(self, func):
         """
         Apply a function to each RDD in this DStream.
         """
+        if func.func_code.co_argcount == 1:
+            old_func = func
+            func = lambda t, rdd: old_func(rdd)
         jfunc = TransformFunction(self._sc, func, self._jrdd_deserializer)
         api = self._ssc._jvm.PythonDStream
         api.callForeachRDD(self._jdstream, jfunc)