如何按指定列唯一值将PySpark DataFrame拆分为DataFrame列表

方案说明 #

要复现pandas中groupby拆分为子DataFrame列表的效果，不需要强制拼接字符串生成唯一键，直接基于多列动态构造过滤条件的方式更稳妥，能避免拼接分隔符与字段值冲突的问题，同时完全支持分组列在运行时动态调整。

实现代码 #

from pyspark.sql import SparkSession
from functools import reduce
import pyspark.sql.functions as F

# 初始化Spark会话
spark_session = SparkSession.builder.enableHiveSupport().getOrCreate()

# 测试数据集
columns = ["language","users_count","apple"]
data = [("Java", 1, 0.0), ("Scala", 4, -4.0), ("Java", 1, 0.0)]
pyspark_df = spark_session.createDataFrame(data).toDF(*columns)

# 可动态修改的分组列列表
column_list = ["language", "users_count"]

# 1. 提取分组列的所有唯一值组合
unique_group_combo = pyspark_df.select(column_list).distinct().collect()

# 2. 遍历组合逐次过滤，生成子DataFrame列表
unique_dfs = []
for combo in unique_group_combo:
    # 动态拼接多列等值过滤条件
    filter_cond = reduce(
        lambda cond1, cond2: cond1 & cond2,
        [F.col(col_name) == combo[col_name] for col_name in column_list]
    )
    sub_df = pyspark_df.filter(filter_cond)
    unique_dfs.append(sub_df)

注意事项 #

• 该方案和pandas groupby输出逻辑一致，每个子DataFrame对应一组分组列的唯一值组合，分组列的数量、名称修改只需要调整column_list变量即可，不需要改动其他逻辑。
• 如果分组的唯一值量级很大（比如超过1000组），不建议使用拆分到Python列表的方案：PySpark是分布式计算引擎，大量小DataFrame会打破分布式执行的优势，带来严重的性能损耗，这类场景优先直接使用groupBy、窗口函数等原生API完成计算，不要把数据拆分到本地列表处理。
• 如果偏好你提到的拼接唯一键的实现思路，注意使用不会出现在业务字段中的特殊分隔符（比如\0）避免键冲突，参考实现如下：

# 生成拼接唯一键
df_with_uid = pyspark_df.withColumn(
    "group_uid",
    F.concat_ws("\0", *[F.col(c).cast("string") for c in column_list])
)
# 拉取所有唯一键
all_uid = [r.group_uid for r in df_with_uid.select("group_uid").distinct().collect()]
# 过滤得到子表后删除额外生成的唯一键列
unique_dfs = [df_with_uid.filter(F.col("group_uid") == uid).drop("group_uid") for uid in all_uid]

内容的提问来源于stack exchange，提问作者Nipun Kumar Goel