如何从Spark DataFrame的层级关系生成指定格式的字典

问题描述 #

我拥有如下Spark DataFrame：

schema = 'EMPLOYEE_NUMBER int, MANAGER_EMPLOYEE_NUMBER int'
employees = spark.createDataFrame(
[[801,None], 
[1016,801], 
[1003,801], 
[1019,801], 
[1010,1003], 
[1004,1003], 
[1001,1003],
[1012,1004], 
[1002,1004], 
[1015,1004], 
[1008,1019], 
[1006,1019], 
[1014,1019],
[1011,1019]], schema=schema)

希望从该DataFrame生成类似{801:[1003,1019,1016], 1019:[1014,1011,1008,1006], 1003:[1010,1001,1004]}的字典，是否可以实现？

解决方案 #

完全可以实现，步骤如下：

• 过滤并分组数据：先过滤掉MANAGER_EMPLOYEE_NUMBER为null的行（这类是顶层管理者，无需作为下属被收集），再按管理者编号分组，收集对应下属的员工编号列表。
• 转换为字典：将分组后的Spark DataFrame转换为Python字典。

具体代码 #

from pyspark.sql import functions as F

# 分组收集下属员工编号列表
grouped_df = employees.filter(employees.MANAGER_EMPLOYEE_NUMBER.isNotNull()) \
    .groupBy("MANAGER_EMPLOYEE_NUMBER") \
    .agg(F.collect_list("EMPLOYEE_NUMBER").alias("subordinates"))

# 转换为Python字典
result_dict = dict(grouped_df.rdd.map(lambda row: (row.MANAGER_EMPLOYEE_NUMBER, row.subordinates)).collect())
print(result_dict)

补充说明 #

• collect_list会保留数据原有顺序，若需要对下属列表排序，可在收集前添加orderBy，或对列表手动排序：

# 排序后的版本
grouped_df = employees.filter(employees.MANAGER_EMPLOYEE_NUMBER.isNotNull()) \
    .orderBy("EMPLOYEE_NUMBER") \
    .groupBy("MANAGER_EMPLOYEE_NUMBER") \
    .agg(F.collect_list("EMPLOYEE_NUMBER").alias("subordinates"))

• 也可通过转换为Pandas DataFrame再生成字典：

import pandas as pd
result_dict = grouped_df.toPandas().set_index("MANAGER_EMPLOYEE_NUMBER")["subordinates"].to_dict()

运行后即可得到符合需求的管理者-下属映射字典。

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