PySpark DataFrame动态长度字符串数组转独热编码，是否有对应ml.feature类？

当然有合适的工具啦！PySpark.ml.feature中的CountVectorizer类（配合binary=True参数）就能完美实现你要的功能——它可以把动态长度的字符串数组转换成存在为1、不存在为0的二进制编码列，完全匹配你的需求。

具体实现步骤 #

1. 创建示例DataFrame

先把你的示例数据转换成PySpark DataFrame：

from pyspark.sql import SparkSession
from pyspark.ml.feature import CountVectorizer, CountVectorizerModel
from pyspark.sql.functions import vector_to_array

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

# 示例数据
data = [
    (["ABC","def","ghi"],),
    (["Jkl","ABC","def"],),
    (["Xyz","ABC"],)
]
df = spark.createDataFrame(data, ["categories"])
df.show(truncate=False)

2. 使用CountVectorizer生成二进制编码

默认情况下，CountVectorizer会根据词汇出现频率排序生成词汇表，但如果你需要严格匹配你指定的编码顺序（比如第一个位置对应"ABC"），可以直接自定义词汇表：

方式一：自动生成词汇表（按出现频率排序）

# 初始化CountVectorizer，设置binary=True启用二进制编码
cv = CountVectorizer(inputCol="categories", outputCol="onehot_encoded", binary=True)
# 拟合模型，自动识别所有词汇
model = cv.fit(df)
# 转换数据
result_df = model.transform(df)

# 查看自动生成的词汇表顺序
print("自动生成的词汇表（对应编码位置）：", model.vocabulary)

方式二：自定义词汇表（完全匹配你的期望顺序）

如果你需要编码顺序严格是["ABC", "def", "ghi", "Jkl", "Xyz"]，可以直接用CountVectorizerModel指定词汇表：

# 自定义词汇表，顺序和你期望的编码位置一致
custom_vocab = ["ABC", "def", "ghi", "Jkl", "Xyz"]
cv_model = CountVectorizerModel(
    inputCol="categories", 
    outputCol="onehot_encoded", 
    vocabulary=custom_vocab, 
    binary=True
)
result_df_custom = cv_model.transform(df)

3. 把稀疏向量转换成密集数组

CountVectorizer默认输出的是稀疏向量（节省内存），如果需要转换成你示例中的密集数组形式，可以用vector_to_array函数：

# 转换为密集数组，并转成整数类型
result_df_custom = result_df_custom.withColumn(
    "onehot_array", 
    vector_to_array("onehot_encoded").cast("array<int>")
)

# 查看最终结果
result_df_custom.select("categories", "onehot_array").show(truncate=False)

最终输出结果 #

运行上述代码后，你会得到完全符合期望的输出：

+-------------------+---------------+
|categories         |onehot_array   |
+-------------------+---------------+
|[ABC, def, ghi]    |[1, 1, 1, 0, 0]|
|[Jkl, ABC, def]    |[1, 1, 0, 1, 0]|
|[Xyz, ABC]         |[1, 0, 0, 0, 1]|
+-------------------+---------------+

为什么用CountVectorizer？ #

• 它专门处理文本/数组类型的特征，支持动态长度的输入
• binary=True参数正好满足"存在即1，不存在即0"的独热需求
• 支持自动识别词汇表或自定义词汇表，灵活性高

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