spark中的嵌套for循环并行化

Question

我正在尝试并行化 spark 中的现有算法（以可扩展的方式）。出于问题的目的，我已对其进行了简化，但它看起来像这样：

for p in all_p:
    all_q = calculate1(p)

    results_for_p = []
    for q in all_q:
        results_for_p.append(calculate2(q))

    save_results(results_for_p)

基本上，我已经将 for 循环嵌套了 2 个我想并行运行的长时间运行的函数。但是，嵌套函数calculate2 的参数的大小取决于每个p。

我的尝试是扁平化输入，以便 calculate2 一起在 all_q 和 all_p 上运行：

rdd = sc.parallelize(all_p)
all_q_each_p = rdd.map(calculate1).collect()

# flatten output to something we can parallelize:
all_q_all_p = []
for all_q in all_q_each_p:
    all_q_all_p.append(all_q)

rdd = sc.parallelize(all_q_all_p)
res = rdd.map(calculate2).collect()

# How to do this?? 
collect_and_save_all_results(res)

如何以一种可以很好扩展的方式编写这个？

Answer 1

这正是flatMap 解决的问题类型。 flatMap 默认调整 rdd 的大小。

代码变得简单多了：

rdd = sc.parallelize(all_p)

rdd.flatMap(calculate1).map(
    lambda args: calculate2(*args)
).collect()