Map Partition Over Chunks in Python with Multiprocessing and Mock

Process data in chunks across multiple CPU cores using multiprocessing Pool.map, and mock the chunk function to test partitioning behavior without heavy computation.

Medium Python 3.9+ Aug 9, 2026 Data pipelines & processing 12 views 0 copies

multiprocessing chunking parallel mocking data-pipeline

Python code

21 lines

Python 3.9+

from multiprocessing import Pool
from unittest.mock import patch, Mock

def process_chunk(chunk):
    return [x * x for x in chunk]

def map_partition_over_chunks(data, chunk_size, process_func=process_chunk):
    chunks = [data[i:i + chunk_size] for i in range(0, len(data), chunk_size)]
    with Pool() as pool:
        results = pool.map(process_func, chunks)
    return [item for sublist in results for item in sublist]

if __name__ == "__main__":
    data = [1, 2, 3, 4, 5, 6]
    mock_func = Mock(side_effect=lambda chunk: [x + 1 for x in chunk])
    with patch("__main__.process_chunk", mock_func):
        mocked_output = map_partition_over_chunks(data, 2, process_func=mock_func)
    real_output = map_partition_over_chunks(data, 2, process_func=process_chunk)
    print(f"Mocked result: {mocked_output}")
    print(f"Real result: {real_output}")
    print(f"Mock called {mock_func.call_count} times with chunks of size 2")

Output

stdout

Mocked result: [2, 3, 4, 5, 6, 7]
Real result: [1, 4, 9, 16, 25, 36]
Mock called 3 times with chunks of size 2

How it works

The map_partition_over_chunks function splits the input list into evenly sized chunks and then uses multiprocessing.Pool.map to apply the processing function to each chunk in parallel. In the if __name__ == "__main__" block, a mock object with a side effect increments each number by 1, and the patch context manager replaces the real process_chunk to verify partitioning. After the mock runs, the real function squares numbers in each chunk, demonstrating how the same code performs real work. The output shows the mock result, the real result, and how many times the mock was called, confirming the chunking logic works as intended.

Common mistakes

Running multiprocessing code on Windows without the `if __name__ == "__main__"` guard causes recursive process creation errors.
Forgetting to pass the mock function as the `process_func` argument, so the patch might not affect the internal call.
Assuming `Pool.map` processes chunks in order; it does preserve order, but results may be slow if the function is not CPU-bound.

Variations

Use `multiprocessing.Pool.imap` for lazy evaluation and memory efficiency with large datasets.
Replace `Pool` with `ThreadPool` from `multiprocessing.pool` for I/O-bound tasks.
Use `concurrent.futures.ProcessPoolExecutor` for a modern approach with similar results.

Real-world use cases

Batch processing large CSV files in a data pipeline, where each chunk is transformed and combined.
Parallelizing data cleaning operations on a large dataset in an ETL job to speed up execution.
Testing partition logic in a distributed processing framework by mocking expensive transformations to verify chunk boundaries.

Map Partition Over Chunks in Python with Multiprocessing and Mock

Python code

Output

How it works

Common mistakes

Variations

Real-world use cases

More from Data pipelines & processing

Tutorials

Quizzes

Python code

Output

How it works

Common mistakes

Variations

Real-world use cases

More from Data pipelines & processing

Keep learning

Tutorials

Quizzes