如何在PySpark中实现单列到多二进制列的映射转换？

需求说明 #

将如下原始两列数据：

+-----------------+
| C1     | C2     |
+--------|--------+
| A      | xxx    |
| B      | yyy    |
| A      | yyy    |
| B      | www    |
| B      | xxx    |
| A      | zzz    |
| A      | xxx    |
| A      | yyy    |
+-----------------+

转换为以C1为行、C2的唯一值为列的二进制标记格式：存在对应值标记为1，不存在标记为0，结果如下：

+--------------------------------------------+
| C1     | www    | xxx    | yyy    | zzz    |
+--------|--------|--------|--------|--------+
| A      |   0    |   1    |   1    |   1    |
| B      |   1    |   1    |   1    |   0    |
+--------------------------------------------+

实现代码 #

1. 初始化环境与创建示例数据

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, max, lit

# 创建SparkSession
spark = SparkSession.builder.appName("BinaryPivotDemo").getOrCreate()

# 模拟原始数据
raw_data = [
    ("A", "xxx"),
    ("B", "yyy"),
    ("A", "yyy"),
    ("B", "www"),
    ("B", "xxx"),
    ("A", "zzz"),
    ("A", "xxx"),
    ("A", "yyy")
]

df = spark.createDataFrame(raw_data, schema=["C1", "C2"])

2. 执行透视转换

核心思路是先去重（避免重复行干扰判断），再按C1分组透视C2列，最后将空值填充为0：

# 去重：保留每个C1-C2的唯一组合
distinct_df = df.distinct()

# 透视转换：将C2的唯一值转为列，存在标记1，空值填充0
result_df = distinct_df.groupBy("C1") \
    .pivot("C2") \
    .agg(max(lit(1))) \
    .na.fill(0)

# 查看结果
result_df.show()

代码说明 #

• distinct()：去除重复的C1-C2组合，因为重复行不影响"是否存在"的判断，同时能提升大数据集下的处理效率。
• groupBy("C1").pivot("C2")：按C1分组，将C2的所有唯一值转换为新的列。
• agg(max(lit(1)))：为每个存在的C1-C2组合标记1，分组后无对应值的位置会生成null。
• na.fill(0)：将所有null替换为0，得到最终的二进制标记效果。

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