如何在PySpark中执行OHE并扁平化数据集，使每个Id仅保留一行？

实现思路 #

和你用Pandas的思路一致：先对category列做独热编码，再按id分组聚合取最大值（保留每个类别是否在该id下出现过的标记），最后将聚合后的独热向量拆分为单独列，实现一行对应一个id的格式。

完整代码实现 #

from pyspark.sql import SparkSession
from pyspark.ml.feature import StringIndexer, OneHotEncoder
from pyspark.sql.functions import col, max, udf
from pyspark.ml.linalg import VectorUDT

# 初始化Spark会话
spark = SparkSession.builder.appName("OHE_Flatten").getOrCreate()

# 创建示例输入数据集
data = [(0, "a"), (1, "b"), (2, "c"), (1, "a"), (2, "a"), (0, "c")]
df = spark.createDataFrame(data, ["id", "category"])

# 1. 将类别字符串转为数值索引（独热编码的前置步骤）
indexer = StringIndexer(inputCol="category", outputCol="category_idx")
indexed_df = indexer.fit(df).transform(df)

# 2. 对索引列执行独热编码，输出向量列
encoder = OneHotEncoder(inputCol="category_idx", outputCol="category_ohe")
encoded_df = encoder.fit(indexed_df).transform(indexed_df)

# 3. 按id分组，对独热向量取最大值——同一id下只要出现过某类别，标记就为1
grouped_df = encoded_df.groupBy("id").agg(max("category_ohe").alias("category_ohe"))

# 获取类别标签，用于生成最终列名
category_labels = indexer.fit(df).labels
column_names = [f"category_{label}" for label in category_labels]

# 定义UDF将独热向量转为数组，方便拆分列
vector_to_array = udf(lambda vec: vec.toArray().tolist(), "array<double>")
array_df = grouped_df.withColumn("category_array", vector_to_array(col("category_ohe")))

# 拆分数组为单独列，并转为整数类型
final_df = array_df.select(
    "id",
    *[col("category_array")[i].cast("int").alias(column_names[i]) for i in range(len(column_names))]
)

# 调整列顺序以匹配示例输出（可选）
final_df = final_df.select("id", "category_a", "category_c", "category_b")
final_df.show()

关键步骤说明 #

• StringIndexer：独热编码要求输入为数值型，因此先将字符串类别转为连续索引值。
• OneHotEncoder：将索引列转为稀疏/稠密向量，每个类别对应向量中的一个位置。
• 分组取max：同一id对应的多个独热向量取最大值，能保留所有出现过的类别标记（出现过为1，未出现为0）。
• 向量转列：通过UDF把向量转为数组，再拆分出单独列，实现扁平化效果。

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