Split lists in a column into one hot encoded features in pyspark

user_6383910 0
Viewed 281

I have a pyspark dataframe like the following:

id tmp_list other features
1 ['Spain', 'Italy'] xxx
2 ['Spain', 'France', 'USA', 'India'] yyy
3 ['Spain', 'Germany'] zzz

and a list of countries like the following:

EU_countries = ['Spain', 'Italy', 'France', 'Germany']

I want to do the following:

  1. Extract all unique values from the column tmp_list
  2. Create new columns for all the values present in EU_countries. For the values that don't exist in EU_countries, create a column called other_countries. Essentially, create column for every entry in EU_countries list + an extra column called other_countries.
  3. If an id contains any of the countries from the EU_countries list, the new column Spain should have 1 as value else 0. Same applies to the other countries from EU_countries list.
  4. If an id contains any country that doesn't exist in the EU_countries list, other_countries column should be filled 1 else 0.

This is the final output I am looking for:

id Spain Italy France Germany other_countries other features
1 1 1 0 0 0 xxx
2 1 0 1 0 1 yyy
3 1 0 0 1 0 zzz

I am breaking my head over this. Can somebody please help me with this?

Any help is greatly appreciated! Thank you so much!

pyspark python
2 Answers

Step 1: replace all non-EU countries in tmp_list with a constant string:

from pyspark.sql import functions as F

df = ...
EU_countries = ['Spain', 'Italy', 'France', 'Germany']

def replaceNonEU(c):
    cond = c == EU_countries[0]
    for country in EU_countries[1:]:
        cond |= (c == country)
    return F.when(cond, c).otherwise(F.lit("other_countries"))

df = df.withColumn("tmp_list", F.array_distinct(F.transform("tmp_list", replaceNonEU)))

#+---+--------------------------------+--------------+
#|id |tmp_list                        |other features|
#+---+--------------------------------+--------------+
#|1  |[Spain, Italy]                  |xxx           |
#|2  |[Spain, France, other_countries]|yyy           |
#|3  |[Spain, Germany]                |zzz           |
#+---+--------------------------------+--------------+

Step 2: for each possible value in tmp_list create a new column indicating if the value is an element of tmp_list:

for c in EU_countries + ['other_countries']:
    df = df.withColumn(c, F.array_contains("tmp_list", c).cast("int"))
df = df.drop("tmp_list")

#+---+--------------+-----+-----+------+-------+---------------+
#| id|other features|Spain|Italy|France|Germany|other_countries|
#+---+--------------+-----+-----+------+-------+---------------+
#|  1|           xxx|    1|    1|     0|      0|              0|
#|  2|           yyy|    1|    0|     1|      0|              1|
#|  3|           zzz|    1|    0|     0|      1|              0|
#+---+--------------+-----+-----+------+-------+---------------+
user_21298010

Reasoning and working like I would in pandas.

  1. Explode
  2. Create Category where non EU_countries are attributed as other_countries
  3. get_dummies. On this one I give credit to this post

Code below;

df=df.select('*').withColumn('tmp_list1', F.explode(col('tmp_list')))#Create new column with exploded list
df=df.select('*').withColumn('Cat', when(col('tmp_list1').isin(EU_countries),df.tmp_list1).otherwise('other_countries'))#Create another column Cat
df.groupBy("tmp_list",'other features').pivot("Cat").agg(F.lit(1)).na.fill(0).show()#Get dummies


+---------------------------+--------------+------+-------+-----+-----+---------------+
|tmp_list                   |other features|France|Germany|Italy|Spain|other_countries|
+---------------------------+--------------+------+-------+-----+-----+---------------+
|[Spain, Germany]           |zzz           |0     |1      |0    |1    |0              |
|[Spain, Italy]             |xxx           |0     |0      |1    |1    |0              |
|[Spain, France, USA, India]|yyy           |1     |0      |0    |1    |1              |
+---------------------------+--------------+------+-------+-----+-----+---------------+
user_89869750
Related