Convert a multi-label column to multiple columns in pyspark?

user_7398259 0
Viewed 383

I have a pyspark data frame like so-

enter image description here

This is a simplified version of the problem that I am trying to solve. In fact the "label" column may have as many as 20 items (which will be strings) in one row. Right now we just have 3 options to choose from, i.e. 0, 1, and 2.

For the problem shown in the image, what I want is three more columns - label_0, label_1, and label_2. For example, the table will look like this after the transformation.

enter image description here

This might look similar to doing one-hot encoding. I am finding a hard time doing this in pyspark.

pyspark-dataframes pyspark apache-spark-sql python
1 Answers

For Spark2.4+, you can try this.

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|
#+---+------+------+------+
user_101073890
Related