Add mapPartitionsWithSplit() to PySpark.

1 files changed, 3 insertions, 1 deletions

diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 9f6b507dbd..3d792bbaa2 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -21,6 +21,7 @@ def load_obj():
 
 
 def main():
+    split_index = read_int(sys.stdin)
     num_broadcast_variables = read_int(sys.stdin)
     for _ in range(num_broadcast_variables):
         bid = read_long(sys.stdin)
@@ -32,7 +33,8 @@ def main():
         dumps = lambda x: x
     else:
         dumps = dump_pickle
-    for obj in func(read_from_pickle_file(sys.stdin)):
+    iterator = read_from_pickle_file(sys.stdin)
+    for obj in func(split_index, iterator):
         write_with_length(dumps(obj), old_stdout)