如何遍历分组后的Spark DataFrame并将行数据传入自定义函数

现有写法的性能问题 #

你目前通过df.filter逐用户生成users_list的写法存在明显瓶颈：每执行一次df.filter(F.col("user")==user)都会触发一次全量数据的扫描与任务调度，当用户量级达到十万级以上时，重复扫描全表的开销会快速累积，完全无法满足生产环境的性能要求。
同时需要注意：GroupedData对象本身是懒执行的聚合逻辑封装，无法直接遍历拿到每个分组的明细DataFrame，这也是分组处理场景下非常常见的误区。

分场景的最优实现 #

场景1：自定义函数可基于批量化逻辑改写（优先推荐） #

如果你的三个自定义函数本质是对分组内的行做计算、过滤、特征衍生类操作，完全不需要逐行调用Spark算子，最推荐用applyInPandas实现分组级别的批处理，全程只需要一次全表扫描就能完成所有用户的计算，性能最高：

import pandas as pd
from pyspark.sql import functions as F

# 定义跨所有分组统一的处理逻辑：输入是单个user分组的pandas格式数据，输出是处理后的同结构数据
def batch_process_group(pdf: pd.DataFrame) -> pd.DataFrame:
    # 这里拿到的pdf就是单个user下的全量行，直接把分组数据传入你的自定义函数
    pdf = first_function(pdf)
    pdf = second_function(pdf)
    pdf = third_function(pdf)
    return pdf

# 按user分组，一次调度完成所有分组的处理
result_df = df.groupby("user").applyInPandas(
    batch_process_group,
    schema=df.schema  # 若函数输出的列和原表结构一致直接复用，不一致就按实际返回列定义schema字符串
)

这种方式完全避免了逐用户触发任务的问题，Executor本地直接处理对应分区内的分组数据，不需要把数据拉到Driver侧，是Spark官方推荐的分组自定义处理模式。

场景2：自定义函数必须依赖Spark DataFrame API实现 #

如果你的三个自定义函数内部必须调用Spark原生算子（比如关联维表、执行SQL、写临时表），没法改写成pandas逻辑，那可以通过一次重分区把相同user的数据路由到同一个分区，从根源上降低后续逐用户过滤的扫描开销：

# 仅做一次重分区，将同一user的所有行分配到同一个分区，后续filter操作不会触发全表扫描
df_partitioned = df.repartition(200, "user")  # 分区数根据集群规模调整，默认200即可

# 仅扫描一次拿到所有不重复的user列表
distinct_users = [row["user"] for row in df_partitioned.select("user").distinct().collect()]

for user in distinct_users:
    # 因为同user的数据已经在固定分区，filter只会读取本地分区的对应数据，不会扫全表
    user_detail_df = df_partitioned.filter(F.col("user") == user)
    # 依次将分组数据传入自定义函数
    processed = first_function(user_detail_df)
    processed = second_function(processed)
    processed = third_function(processed)
    # 处理完成后直接落地，不要把所有结果DataFrame存在内存中，避免Driver内存溢出
    processed.write.parquet(f"/tmp/processed_result/{user}", mode="overwrite")

这种写法和你原来的逻辑流程完全兼容，但是通过一次重分区把多次全表扫描的开销降到了仅需一次，性能提升幅度通常在几十到上百倍。

避坑提示 #

• 不要把所有用户处理后的DataFrame存在Python列表中，每个DataFrame都是懒执行的对象，攒在内存里很快就会把Driver的内存占满，处理完一个就落地一个。
• 如果自定义函数的逻辑没有复杂到必须逐用户处理，优先用场景1的applyInPandas方案，性能远高于逐用户遍历的模式。
• 不要尝试直接遍历groupBy返回的GroupedData对象，这类对象不存储实际的分组明细数据，遍历无法拿到单个用户的DataFrame。

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