PySpark分组后展平DataFrame行的实现求助

核心思路 #

通过groupBy按member_id分组，使用collect_list将每组的所有字段打包成数组，最后将Spark DataFrame结果转换为Python原生嵌套列表。

完整代码实现 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import collect_list, array

# 初始化SparkSession
spark = SparkSession.builder.appName("GroupAndCollect").getOrCreate()

# 创建示例DataFrame
data = [
    (1, 35, "Male", "2023-10-01", 200),
    (1, 35, "Male", "2023-10-02", 210),
    (2, 28, "Female", "2023-10-01", 150),
    (2, 28, "Female", "2023-10-02", 160)
]
columns = ["member_id", "age", "gender", "date", "cost"]
df = spark.createDataFrame(data, columns)

# 分组并聚合：将每行的所有字段转为数组，再按member_id收集成列表
grouped_df = df.groupBy("member_id").agg(
    collect_list(array(*columns)).alias("grouped_data")
)

# 转换为Python原生嵌套列表
result = [row.grouped_data for row in grouped_df.collect()]

# 打印结果
print(result)

输出结果 #

运行后会得到你期望的格式：

[
    [[1, 35, 'Male', '2023-10-01', 200], [1, 35, 'Male', '2023-10-02', 210]],
    [[2, 28, 'Female', '2023-10-01', 150], [2, 28, 'Female', '2023-10-02', 160]]
]

关键说明 #

• array(*columns)：把当前行的所有字段打包成一个数组，对应你期望的每个子列表。
• collect_list(...)：按member_id分组后，把每组的所有行数组收集成一个大列表，对应每组的嵌套结构。
• grouped_df.collect()：将Spark分布式数据集拉取到本地，再通过列表推导式提取每组的grouped_data字段，得到最终的嵌套列表。

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