如何为PySpark DataFrame中含多值的列生成哑变量

问题分析 #

你尝试的df.groupby('ID').any().astype(int)无效的原因是：该操作仅对原P1、P2列本身做布尔判断（是否存在非空值），并没有将列中的取值扩展为新的独热编码列，因此无法得到期望的宽表结构。

可行解决方案 #

方法一：使用get_dummies生成独热编码后分组求和 #

先对P1和P2分别生成带前缀的独热编码列，再按ID分组求和（同一ID下出现过的取值会被标记为1，多次出现仍为1）：

import pandas as pd

# 构造原始DataFrame
data = {
    'ID': [0, 0, 1, 1, 2, 3, 3],
    'P1': [447, 448, 447, 450, 450, 451, 452],
    'P2': ['O1', 'O2', 'O2', 'O3', 'O3', 'O4', 'O5']
}
df = pd.DataFrame(data)

# 生成P1、P2的独热编码，添加列前缀
dummies_p1 = pd.get_dummies(df['P1'], prefix='P1')
dummies_p2 = pd.get_dummies(df['P2'], prefix='P2')

# 合并ID与独热编码列
df_encoded = pd.concat([df['ID'], dummies_p1, dummies_p2], axis=1)

# 按ID分组求和，得到最终结果
result = df_encoded.groupby('ID').sum().reset_index()
print(result)

方法二：使用pivot_table分别处理两列后合并 #

通过透视表将P1、P2的取值转为列，再合并结果：

import pandas as pd

# 构造原始DataFrame
data = {
    'ID': [0, 0, 1, 1, 2, 3, 3],
    'P1': [447, 448, 447, 450, 450, 451, 452],
    'P2': ['O1', 'O2', 'O2', 'O3', 'O3', 'O4', 'O5']
}
df = pd.DataFrame(data)

# 处理P1列生成透视表，重命名列名
pivot_p1 = df.pivot_table(
    index='ID', 
    columns='P1', 
    aggfunc=lambda x: 1, 
    fill_value=0
).rename(columns=lambda x: f'P1_{x}')

# 处理P2列生成透视表，重命名列名
pivot_p2 = df.pivot_table(
    index='ID', 
    columns='P2', 
    aggfunc=lambda x: 1, 
    fill_value=0
).rename(columns=lambda x: f'P2_{x}')

# 合并两个透视表结果
result = pd.concat([pivot_p1, pivot_p2], axis=1).reset_index()
print(result)

两种方法都能生成你期望的DataFrame结构，输出结果如下：

ID  P1_447  P1_448  P1_450  P1_451  P1_452  P2_O1  P2_O2  P2_O3  P2_O4  P2_O5
0   0       1       1       0       0       0      1      1      0      0      0
1   1       1       0       1       0       0      0      1      1      0      0
2   2       0       0       1       0       0      0      0      1      0      0
3   3       0       0       0       1       1      0      0      0      1      1

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