python - 在 pyspark 中将多标签列转换为多列？

Question

我有一个像这样的 pyspark 数据框-

这是我试图解决的问题的简化版本。事实上，“标签”列一行中可能有多达 20 个项目(将是字符串)。现在我们只有 3 个选项可供选择，即 0、1 和 2。

对于图中所示的问题，我想要的是另外三列 - label_0、label_1 和 label_2。例如，转换后的表格将如下所示。

这可能看起来类似于进行 one-hot 编码。我发现在 pyspark 中很难做到这一点。

Answer 1

对于Spark2.4+，您可以尝试此操作。

labels=['0','1','2']

from pyspark.sql import functions as F
df.withColumn("struct", F.struct(*[(F.struct(F.expr("size(filter(label,x->x={}))"\
                                                    .format("'"+y+"'"))).alias(y)) for y in labels]))\
            .select("id",*[F.col("struct.{}.col1".format(x)).alias('label'+x) for x in labels]).show()

#+---+------+------+------+
#| id|label0|label1|label2|
#+---+------+------+------+
#|  0|     0|     1|     1|
#|  1|     1|     1|     0|
#|  2|     1|     1|     0|
#|  3|     1|     1|     0|
#+---+------+------+------+