如何将DataFrame中的数组字段赋值给变量？collect()函数失效

假设你使用的是PySpark，以下几种方法可以实现需求：

方法1：通过RDD转为字典（推荐） #

利用collectAsMap()直接将DataFrame转为以Language为键、Department数组为值的字典，再直接赋值变量：

# 将DataFrame转换为键值对字典
dept_map = df.rdd.collectAsMap()

# 直接通过键获取对应数组赋值给变量
English = dept_map.get("English", [])
Spanish = dept_map.get("Spanish", [])

方法2：转为Pandas DataFrame处理 #

如果数据量不大，可以转成Pandas DataFrame后用字典映射：

# 转换为Pandas DataFrame
pdf = df.toPandas()
# 构建Language到Department的映射字典
dept_dict = pdf.set_index("Language")["Department"].to_dict()

# 赋值变量
English = dept_dict.get("English", [])
Spanish = dept_dict.get("Spanish", [])

方法3：使用collect()手动解析 #

如果坚持用collect()，需要遍历返回的Row对象提取数据：

# 收集所有行数据
rows = df.collect()

# 初始化变量
English = []
Spanish = []

# 遍历行提取对应值
for row in rows:
    if row.Language == "English":
        English = row.Department
    elif row.Language == "Spanish":
        Spanish = row.Department

注意：如果DataFrame中存在重复的Language值，方法1和方法2会保留最后一条对应的数据，方法3同理，若需处理重复值需额外逻辑。

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