如何在PySpark中并行执行列表上的重复操作？

程序中存在这样的场景：一个函数接收列表，需要对列表中的每个元素重复执行操作。以统计PySpark DataFrame满足条件的行为例，当前实现代码如下：

l = [2, 4, 5]
res = []

for x in l:
    val = df.where((col('id') == x) | (col('id2') == x)).count()
    res.append(val)

现在想让多个worker同时计算每个val，也就是让每个worker独立计算对应元素的结果并汇总到结果列表中。尝试过用foreach，但因为遍历对象是普通列表而非RDD或DataFrame，无法直接使用该方法。

这里提供两种可行的并行计算方案：

方案一：将列表转为RDD实现分布式计算 #

把普通列表转换成RDD，就能利用Spark的分布式能力让worker并行处理每个元素：

from pyspark.sql import SparkSession
from pyspark.sql.functions import col

# 初始化SparkSession（如果未初始化）
spark = SparkSession.builder.appName("ParallelCount").getOrCreate()

l = [2, 4, 5]
# 将列表转为RDD
rdd = spark.sparkContext.parallelize(l)

# 定义每个元素的计算逻辑
def count_matches(x):
    return df.where((col('id') == x) | (col('id2') == x)).count()

# 并行执行计算，collect()获取结果列表
res = rdd.map(count_matches).collect()

注意：要确保df能被序列化传递到worker节点，保证DataFrame的元数据可在集群中访问。

方案二：使用Python多进程（适合小规模场景） #

如果数据规模不大，也可以用Python标准库的multiprocessing实现本地并行，不过这种方式不会用到Spark的worker，而是利用本地CPU核心：

from multiprocessing import Pool
from pyspark.sql.functions import col

l = [2, 4, 5]

# 注意：多进程环境中需确保df可正确传递，或在子进程中重新初始化Spark连接
def count_matches(x):
    return df.where((col('id') == x) | (col('id2') == x)).count()

# 创建进程池并行计算
with Pool(processes=3) as pool:
    res = pool.map(count_matches, l)

这种方式局限于本地资源，无法利用Spark集群的分布式能力，仅适合小批量数据场景。

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