如何在PySpark中并行化for循环，优化Pandas迁移代码？

原有代码问题 #

• 三重for循环（遍历机构、遍历用户、遍历日期）全部在Driver端执行，完全没有利用Spark分布式计算能力
• 频繁调用collect()、filter()拉取数据和触发小作业，性能极低，数据量稍大就会出现Driver OOM或者任务超时

最优实现思路 #

全程使用Spark原生算子执行，全链路分布式计算，没有任何Driver端遍历业务数据的循环：

1. 先对全表按机构、用户、日期去重，避免重复计算
2. 用窗口函数按机构+用户分组后对日期排序，取每个日期的前一个相邻日期
3. 计算相邻日期间隔、当前日期所属月份
4. 按机构、月份、间隔天数分组统计出现次数
5. 最后仅将极小的统计结果拉取到Driver，转换为所需嵌套字典格式

完整代码实现 #

from pyspark.sql import functions as F
from pyspark.sql.window import Window
from collections import defaultdict

def execute():
    df = load_data() # 输入Spark DataFrame，对应schema：institution_id、user_id、st_date
    # 若实际机构列名为inst_id，可按需替换下面的institution_id字段名

    # 1. 去重+统一转换st_date为日期类型
    dedup_df = df.dropDuplicates(["institution_id", "user_id", "st_date"]) \
                 .withColumn("st_date", F.to_date("st_date"))

    # 2. 定义窗口规则：按机构+用户分区，按日期升序排序
    w = Window.partitionBy("institution_id", "user_id").orderBy("st_date")

    # 3. 计算相邻日期差值、当前日期所属月份
    date_diff_df = dedup_df.withColumn("prev_date", F.lag("st_date", 1).over(w)) \
                           .filter(F.col("prev_date").isNotNull()) # 过滤每个用户的第一条无前置日期的记录
                           .withColumn("cycle_days", F.datediff("st_date", "prev_date")) \
                           .withColumn("month", F.month("st_date"))

    # 4. 按维度分组统计次数
    stat_df = date_diff_df.groupBy("institution_id", "month", "cycle_days") \
                          .agg(F.count("*").alias("cnt"))

    # 5. 转换为要求的嵌套字典格式：{机构id: {月份: {间隔天数: 计数}}}
    monthly_distributions = defaultdict(lambda: defaultdict(lambda: defaultdict(int)))
    for row in stat_df.collect():
        inst_id = row["institution_id"]
        month = row["month"]
        cycle_days = row["cycle_days"]
        cnt = row["cnt"]
        monthly_distributions[inst_id][month][cycle_days] = cnt

    # 按需可转为普通字典返回
    # monthly_distributions = {k: {k2: dict(v2) for k2, v2 in v.items()} for k, v in monthly_distributions.items()}
    return monthly_distributions

性能优势 #

• 所有计算逻辑全部在Executor端分布式并行执行，适配EMR集群的资源调度能力，数据量越大性能优势越明显
• 仅最后一次collect拉取最终统计结果，数据量极小，不会出现Driver内存溢出问题

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