数据帧列中动态长度的 Pyspark 字符串数组进行 onehot 编码

我想转换包含以下字符串的列：

 ["ABC","def","ghi"] 
 ["Jkl","ABC","def"]
 ["Xyz","ABC"]

进入这样的编码列：

 [1,1,1,0,0]
 [1,1,0,1,0]
 [0,1,0,0,1]

pyspark.ml.feature 中有一个类吗？

编辑：在编码列中，第一个条目始终对应于值“ABC”等。1 表示“ABC”存在，而 0 表示它不存在于相应的行中。

你也许可以使用计数向量化器 http://spark.apache.org/docs/2.4.0/api/python/pyspark.ml.html#pyspark.ml.feature.CountVectorizer，下面是一个例子：

Update:删除了删除数组中重复项的步骤，您可以设置binary=True设置 CountVectorizer 时：

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

df = spark.createDataFrame([
        (["ABC","def","ghi"],)
      , (["Jkl","ABC","def"],)
      , (["Xyz","ABC"],)
    ], ['arr']
)

创建 CountVectorizer 模型：

cv = CountVectorizer(inputCol='arr', outputCol='c1', binary=True)

model = cv.fit(df)

vocabulary = model.vocabulary
# [u'ABC', u'def', u'Xyz', u'ghi', u'Jkl']

创建 UDF 将向量转换为数组

udf_to_array = udf(lambda v: v.toArray().tolist(), 'array<double>')

获取向量并检查内容：

df1 = model.transform(df)

df1.withColumn('c2', udf_to_array('c1')) \
   .select('*', *[ col('c2')[i].astype('int').alias(vocabulary[i]) for i in range(len(vocabulary))]) \
   .show(3,0)
+---------------+-------------------------+-------------------------+---+---+---+---+---+
|arr            |c1                       |c2                       |ABC|def|Xyz|ghi|Jkl|
+---------------+-------------------------+-------------------------+---+---+---+---+---+
|[ABC, def, ghi]|(5,[0,1,3],[1.0,1.0,1.0])|[1.0, 1.0, 0.0, 1.0, 0.0]|1  |1  |0  |1  |0  |
|[Jkl, ABC, def]|(5,[0,1,4],[1.0,1.0,1.0])|[1.0, 1.0, 0.0, 0.0, 1.0]|1  |1  |0  |0  |1  |
|[Xyz, ABC]     |(5,[0,2],[1.0,1.0])      |[1.0, 0.0, 1.0, 0.0, 0.0]|1  |0  |1  |0  |0  |
+---------------+-------------------------+-------------------------+---+---+---+---+---+