如何在PySpark中展开数组列生成布尔值列

给定如下Pandas DataFrame（Answers列为数组类型）：

+------------+-----------------+------------------------------------+
| Name       |   Age           | Answers                            |
+------------+-----------------+------------------------------------+
| Maria      | 23              | [apple, mango, orange, banana]     | 
| John       | 55              | [apple, orange, banana]            |
| Brad       | 44              | [banana]                           |
| Alex       | 55              | [apple, mango, orange, banana]     |
+------------+-----------------+------------------------------------+

需要生成如下格式的输出：将所有出现过的answer选项展开为独立行，并为每个用户标记是否包含该选项（value列，存在为True，不存在为False）：

+-----+---+--------+-------+                                                               
| Name|Age|  answer| value |
+-----+---+--------+-------+
|Maria| 23|   apple| True  |
|Maria| 23|   mango| True  |
|Maria| 23|  orange| True  |
|Maria| 23|  banana| True  |
| John| 55|   apple| True  |
| John| 55|   mango| False |
| John| 55|  orange| True  |
| John| 55|  banana| True  |
| Brad| 44|   apple| False |
| Brad| 44|   mango| False |
| Brad| 44|  orange| False |
| Brad| 44|  banana| True  |
|Alex | 55|   apple| True  |
|Alex | 55|   mango| True  |
|Alex | 55|  orange| True  |
|Alex | 55|  banana| True  |
+-----+---+--------+-------+

解决方案 #

方法一：笛卡尔积+行判断 #

适合数据量较小的场景，逻辑直观：

1. 提取所有唯一选项
   先从Answers列的所有数组中，收集并去重得到全部可选值：

   import pandas as pd
   
   # 构造原始DataFrame
   df = pd.DataFrame({
       'Name': ['Maria', 'John', 'Brad', 'Alex'],
       'Age': [23, 55, 44, 55],
       'Answers': [
           ['apple', 'mango', 'orange', 'banana'],
           ['apple', 'orange', 'banana'],
           ['banana'],
           ['apple', 'mango', 'orange', 'banana']
       ]
   })
   
   # 获取所有唯一的answer选项并排序
   all_answers = sorted({item for sublist in df['Answers'] for item in sublist})
   

2. 生成用户与选项的笛卡尔积
   通过添加临时键实现原DataFrame与选项列表的交叉连接，得到每个用户对应所有选项的行：

   # 创建选项DataFrame
   answers_df = pd.DataFrame({'answer': all_answers})
   # 交叉连接生成全量组合
   cross_df = df.assign(key=1).merge(answers_df.assign(key=1), on='key').drop('key', axis=1)
   

3. 标记选项是否存在
   用apply判断每个选项是否在用户的Answers数组中，生成value列：

   cross_df['value'] = cross_df.apply(lambda row: row['answer'] in row['Answers'], axis=1)
   # 调整列顺序得到最终结果
   result = cross_df[['Name', 'Age', 'answer', 'value']]
   print(result)
   

方法二：Explode+透视表补全 #

适合数据量较大的场景，避免apply的性能损耗：

import pandas as pd

# 构造原始DataFrame
df = pd.DataFrame({
    'Name': ['Maria', 'John', 'Brad', 'Alex'],
    'Age': [23, 55, 44, 55],
    'Answers': [
        ['apple', 'mango', 'orange', 'banana'],
        ['apple', 'orange', 'banana'],
        ['banana'],
        ['apple', 'mango', 'orange', 'banana']
    ]
})

# 展开Answers列，标记存在的选项为True
exploded = df.explode('Answers').rename(columns={'Answers': 'answer'})
exploded['value'] = True

# 用透视表补全缺失的选项，填充False后重新展开
result = exploded.pivot_table(
    index=['Name', 'Age'],
    columns='answer',
    values='value',
    fill_value=False
).stack().reset_index(name='value')

# 调整列顺序
result = result[['Name', 'Age', 'answer', 'value']]
print(result)

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