如何对DataFrame行分组并对每组行执行操作？

一种显而易见（但效率低下）的方法是收集组ID并逐个遍历：

vals = df.groupBy('SomeField').agg(F.count("*").alias("Count")).collect()
for val in vals:
    group_df = df.where(df.SomeField == val.SomeField)
    do_something(group_df)

有没有更优的实现方式？

原方法存在两个核心效率问题：

• collect()会将分组统计结果拉取到Driver端，分组数量较大时易引发Driver内存溢出
• 每次循环调用df.where()都会重新扫描全量原始数据，重复计算成本极高

针对不同业务场景，推荐以下高效实现方式：

1. 原生Spark分组算子：foreachGroup/mapGroups #

如果仅需对每个分组执行侧处理（如写入存储、调用外部接口），foreachGroup是最直接的选择——它会在Executor端直接处理分组数据，无需拉取到Driver：

def process_group(key, group_df):
    # key为当前分组的SomeField值，group_df是该分组的DataFrame
    do_something(group_df)

df.groupBy("SomeField").foreachGroup(process_group)

如果需要对分组数据做转换并返回结果，可使用mapGroups（需基于RDD操作）：

from pyspark.sql import Row

def transform_group(key, rows_iter):
    group_rows = list(rows_iter)
    # 执行自定义转换逻辑
    transformed_rows = [Row(SomeField=key, processed_col=row.col * 2) for row in group_rows]
    return transformed_rows

processed_rdd = df.rdd.groupBy(lambda row: row.SomeField).flatMap(lambda x: transform_group(x[0], x[1]))
processed_df = processed_rdd.toDF(df.schema)

2. Pandas UDF（适合依赖Pandas生态的场景） #

若do_something需要用到Pandas的数据分析能力，使用Spark Pandas UDF能兼顾分布式性能与Pandas的便捷性：

from pyspark.sql.functions import pandas_udf, PandasUDFType

@pandas_udf(df.schema, PandasUDFType.GROUPED_MAP)
def process_with_pandas(group_df):
    # group_df为当前分组的Pandas DataFrame
    do_something(group_df)
    return group_df  # 返回处理后的Pandas DataFrame

processed_df = df.groupBy("SomeField").apply(process_with_pandas)

3. 预分区后处理（适合高频分组场景） #

如果后续会多次按SomeField分组处理，可先对DataFrame按该字段分区，减少重复分区开销：

# 按SomeField重分区，确保同组数据落在同一Executor分区
df_partitioned = df.repartition("SomeField")

def process_partition(rows_iter):
    from collections import defaultdict
    groups = defaultdict(list)
    for row in rows_iter:
        groups[row.SomeField].append(row)
    # 遍历分区内的所有分组进行处理
    for key, group_rows in groups.items():
        do_something(spark.createDataFrame(group_rows, df.schema))

df_partitioned.rdd.foreachPartition(process_partition)

以上方案均避免了原方法的重复扫描与Driver端数据拉取问题，在大数据量场景下能显著提升处理效率。

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