Bug: Serialize all data at once may cause MemoryError
Status: Open | Priority: Critical | Resolution: 
Created: 25/May/24 07:28
Affects Version/s: 2.2.1

In worker.py, there is a function called process(), the iterator loads all data at once
{code:java}
def process():
       iterator = deserializer.load_stream(infile)
       serializer.dump_stream(func(split_index, iterator), outfile){code}
It will cause MemoryError when working on large scale data, For the reason that I have indeed encountered this situation as below:
{code:java}
MemoryError at org.apache.spark.api.python.PythonRunner$$anon$1.read(PythonRDD.scala:203) at org.apache.spark.api.python.PythonRunner$$anon$1.<init>(PythonRDD.scala:244) at org.apache.spark.api.python.PythonRunner.compute(PythonRDD.scala:162) at org.apache.spark.sql.execution.python.BatchEvalPythonExec$$anonfun$doExecute$1.apply(BatchEvalPythonExec.scala:144) at org.apache.spark.sql.execution.python.BatchEvalPythonExec$$anonfun$doExecute$1.apply(BatchEvalPythonExec.scala:87) at org.apache.spark.rdd.RDD$$anonfun$mapPartitions$1$$anonfun$apply$23.apply(RDD.scala:797) at org.apache.spark.rdd.RDD$$anonfun$mapPartitions$1$$anonfun$apply$23.apply(RDD.scala:797) at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:38) at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:323) at org.apache.spark.rdd.RDD.iterator(RDD.scala:287) at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:38) at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:323) at org.apache.spark.rdd.RDD.iterator(RDD.scala:287) at org.apache.spark.scheduler.ResultTask.runTask(ResultTask.scala:89) at org.apache.spark.scheduler.Task.run(Task.scala:109) at org.apache.spark.executor.Executor$TaskRunner$$anon$2.run(Executor.scala:355) at java.security.AccessController.doPrivileged(Native Method) at javax.security.auth.Subject.doAs(Subject.java:422) at org.apache.hadoop.security.UserGroupInformation.doAs(UserGroupInformation.java:1721) at org.apache.spark.executor.Executor$TaskRunner.run(Executor.scala:353) at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1149) at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:624) at java.lang.Thread.run(Thread.java:748) 2024-05-22 16:50:03,173 INFO org.apache.spark.scheduler.TaskSetManager: Starting task 0.1 in stage 5.0 (TID 21, saturndatanode3, executor 2, partition 0, ANY, 5075 bytes) 2024-05-22 16:50:03,174 INFO org.apache.spark.scheduler.TaskSetManager: Lost task 1.0 in stage 5.0 (TID 19) on saturndatanode3, executor 2: org.apache.spark.api.python.PythonException (Traceback (most recent call last): File "xxxxxx/spark/python/lib/pyspark.zip/pyspark/worker.py", line 200, in main process() File "xxxxx/spark/python/lib/pyspark.zip/pyspark/worker.py", line 195, in process serializer.dump_stream(func(split_index, iterator), outfile) File "xxxxx/spark/python/lib/pyspark.zip/pyspark/worker.py", line 106, in <lambda> func = lambda _, it: map(mapper, it) File "<string>", line 1, in <lambda> File "xxxxx/spark/python/lib/pyspark.zip/pyspark/worker.py", line 73, in <lambda> return lambda *a: f(*a){code}
I did some tests by adding memory monitor code, I found that this code takes up a lot of memory during execution：
{code:java}
start_memory = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
print(f"Memory usage at the beginning: {start_memory} KB")iterator = deserializer.load_stream(infile)
serializer.dump_stream(func(split_index, iterator), outfile)end_memory = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
print(f"Memory usage at the end: {end_memory} KB")memory_difference = end_memory - start_memory
print(f"Memory usage changes:{memory_difference} KB"){code}
Can I process the data in the iterator in batches as below？
{code:java}
def process():
    iterator = deserializer.load_stream(infile)
    def batched_func(iterator, func, serializer, outfile):
        batch = []
        count = 0
        for item in iterator:
            batch.append(item)
            count += 1
            // Process the data in the iterator in batches, with 10000 entries each time.
            if count >= 10000:
                serializer.dump_stream(func(split_index, batch), outfile)
                batch = []
                count = 0
            if batch:
                serializer.dump_stream(func(split_index, batch), outfile)
        batched_func(iterator, func, serializer, outfile){code}
I test with code as above， it works well with lower memory each time.