PySpark如何对列表类型标签特征做One-Hot Encoding

我需要预处理供机器学习使用的电视剧数据集，单条记录对应一部电视剧，包含一个存储该剧关联标签的列表类型特征，处理后的特征将用于后续分析。现有两个方案待选：直接对该列表特征做One-Hot编码，或是先拆分提取列表内所有独立标签再做编码。
数据集示例：

当前使用的处理代码：

from pyspark.ml.feature import StringIndexer
from pyspark.ml.feature import OneHotEncoder

indexer = StringIndexer(inputCol="tags", outputCol="tagsIndex")
df = indexer.fit(df).transform(df)

ohe = OneHotEncoder(inputCol="tagsIndex", outputCol="tagsOHEVector")
df = ohe.fit(df).transform(df)

你当前的写法完全不可用，必须先拆分提取所有独立标签，再做多热（Multi-Hot）编码，不能直接对整个标签列表做普通One-Hot编码。

现有代码的问题 #

• StringIndexer会把每一个完整的标签列表当成一个独立的类别值处理：比如两部剧标签分别是["古装","权谋"]和["古装","爱情"]，会被识别为两个毫无关联的独立类别，编码结果完全体现不出两部剧共有的「古装」标签属性，后续分析或模型训练根本无法捕捉标签的共性特征。
• 如果标签列表的存储顺序不固定，哪怕两个剧集标签完全一致只是顺序不同（比如["悬疑","犯罪"]和["犯罪","悬疑"]），也会被识别为不同类别，生成无意义的冗余编码。

标准处理方案 #

多标签类型字段的通用编码方式是多热编码：先提取数据集中所有出现过的独立标签作为单独的特征维度，某条记录包含对应标签则该维度赋值为1，不包含则赋值为0。
在PySpark中不需要手动写explode、pivot逻辑，直接用内置的CountVectorizer即可高效实现，设置binary=True就会输出纯0/1的多热向量，不会统计标签出现次数：

from pyspark.ml.feature import CountVectorizer
from pyspark.sql.functions import split, col

cv = CountVectorizer(inputCol="tags", outputCol="tagsMultiHotVec", binary=True)
cv_model = cv.fit(df)
df = cv_model.transform(df)

# 可通过cv_model.vocabulary获取向量每一维对应的标签名称

注意：如果你的tags列是逗号分隔的字符串类型、而非数组类型，需要先执行拆分逻辑再传入CountVectorizer：

df = df.withColumn("tags", split(col("tags"), ","))

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