如何让multiprocessing.Pool().starmap()返回传入的可迭代参数

不需要修改原函数的返回值，你可以通过以下方式关联输入参数和函数结果：

1. 提前保存参数列表（内存充足时首选） #

先把所有要传入的(i,j)参数对保存成一个列表，之后将这个列表和starmap返回的结果一一配对：

import multiprocessing as mp
import pandas as pd

# 生成所有参数对列表
params = [(i, j) for j in range(y) for i in range(x)]

with mp.Pool() as pool:
    # 传入参数列表并执行函数
    results = pool.starmap(func, params)

# 将参数和结果合并成DataFrame所需的格式
data = [(i, j, res) for (i, j), res in zip(params, results)]
df = pd.DataFrame(data, columns=["i", "j", "k"])

starmap会严格按照传入的可迭代对象的顺序返回结果，因此参数列表和结果列表的索引完全对应，直接zip即可正确配对。

2. 用生成器复制避免内存占用（数据量极大时） #

如果x和y过大，生成参数列表会占用过多内存，可以用itertools.tee复制参数生成器，一个给starmap使用，另一个用来和结果配对：

import multiprocessing as mp
import pandas as pd
from itertools import tee

# 创建参数生成器
param_gen = ((i, j) for j in range(y) for i in range(x))
# 复制生成器，两个生成器会独立迭代
param_gen_for_pool, param_gen_for_pairing = tee(param_gen)

with mp.Pool() as pool:
    results = pool.starmap(func, param_gen_for_pool)

# 配对参数和结果
data = [(i, j, res) for (i, j), res in zip(param_gen_for_pairing, results)]
df = pd.DataFrame(data, columns=["i", "j", "k"])

这种方式不会一次性把所有参数加载到内存，适合处理超大规模的循环任务。

内容的提问来源于stack exchange，提问作者TheRavenSpectre