PySpark实现类似pandas.get_dummies的独热编码方案问询

针对你的大规模PySpark DataFrame，无需转换为Pandas，直接用纯PySpark操作即可实现需求，以下是两种简单的实现方法：

方法一：手动创建标记列（适合已知固定取值的场景） #

已知Value列只有value1和value2两个取值，可直接生成对应布尔标记列，再分组聚合后转整数：

from pyspark.sql.functions import col, max

# 1. 为每个Value取值创建布尔标记列
df_with_flags = df.withColumn("value1", col("Value") == "value1") \
                  .withColumn("value2", col("Value") == "value2")

# 2. 按Key分组，取布尔列的最大值（只要组内存在对应Value就为True）
grouped_df = df_with_flags.groupBy("Key1", "Key2", "Key3") \
                          .agg(
                              max("value1").alias("value1"),
                              max("value2").alias("value2")
                          )

# 3. 将布尔列转换为整数类型（True→1，False→0）
result_df = grouped_df.withColumn("value1", col("value1").cast("integer")) \
                      .withColumn("value2", col("value2").cast("integer"))

方法二：使用Pivot（通用型，适合Value取值不固定的场景） #

如果后续Value可能新增取值，用pivot更灵活，无需修改代码：

from pyspark.sql.functions import lit, max

# 1. 添加一个固定为True的标记列，用于后续聚合
df_flagged = df.withColumn("flag", lit(True))

# 2. 按Key分组，以Value为列进行pivot，聚合取max（存在则为True）
pivoted_df = df_flagged.groupBy("Key1", "Key2", "Key3") \
                       .pivot("Value") \
                       .agg(max("flag")) \
                       .fillna(False)  # 填充不存在的取值为False

# 3. 将布尔列转换为整数类型
result_df = pivoted_df.withColumn("value1", col("value1").cast("integer")) \
                      .withColumn("value2", col("value2").cast("integer"))

关键说明 #

• 两种方法均采用布尔值取max的聚合逻辑，确保最终是1/0的独热编码（而非计数）：只要分组内存在对应Value，结果就是True（转整数为1），否则为False（转整数为0）。
• 全程基于PySpark分布式处理，不会出现大规模数据转Pandas导致的内存溢出问题。

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