best way to perform one-hot encoding and explode column with prefix pyspark

user_10970202 0
Viewed 277

Trying to replicate pandas code in pyspark 2.x.

say I have dataframe as follows:

    age      education    country
0   22          A          Canada
1   34          B          Mongolia
2   55          A          Peru
3   44          C          Korea

Usually in pandas I would scale numerical columns and one hot encode categorical and get:

    age      education_A education_B education_C    country_Canada   country_Mongolia ...
0   0            1           0           0                1                 0 
1   0.3          0           1           0                0                 0
2   1            1           0           0                0                 0         ...
3   0.7          0           0           1                ...              ...

In pyspark I've done

str_indexer1 = StringIndexer(inputCol="education", outputCol=education+"_si", handleInvalid="skip")
str_indexer2 = StringIndexer(inputCol="country", outputCol=country+"_si", handleInvalid="skip")
mod_df = str_indexer1.fit(df).transform(df)
mod_df = str_indexer2.fit(df).transform(mod_df)

ohe1 = OneHotEncoder(inputCol="education", outputCol=education+"_ohe")
ohe2 = OneHotEncoder(inputCol="country", outputCol=country+"_ohe")

ohe1.fit(mod_df).transform(mod_df)

This gives me

    age      education   country          education_si    country_si     education_ohe
0   0            A         Canada              1              1           (1,0,0,0)
1   0.3          B         Mongolia            2              2           (0,1,0,0)
2   1            A         Peru                1              3           (1,0,0,0)
3   0.7          C         Korea               3              4           (0,0,1,0)

From here I cannot find out how to explode education_ohe into education_A, etc...

How can I do this and also is more efficient way to perform ohe and scaler in large dataframe?

pyspark pandas apache-spark python
1 Answers

There is a built in oneHotEncoder in pyspark's functions, but I could not get it to provide true one-hot encoded columns.

That being said the following code will get the desired result.

Using the following dataframe.

df.show()
+---+---------+--------+
|age|education| country|
+---+---------+--------+
| 22|        A|  Canada|
| 34|        B|Mongolia|
| 55|        A|    Peru|
| 44|        C|   Korea|
+---+---------+--------+

You can select all the distinct column values, and iteratively create additional columns.

variables_dict = {}

for col in df.columns:
    variables_dict[col] = [row[col] for row in df.distinct().collect()]

for col in variables_dict.keys():
    for val in variables_dict[col]:
        df = df.withColumn("{}_{}".format(col, val), functions.when((df[col] == val), 1).otherwise(0))

df.show()
+---+---------+--------+------+------+------+------+-----------+-----------+-----------+--------------+-------------+------------+----------------+
|age|education| country|age_22|age_44|age_55|age_34|education_A|education_C|education_B|country_Canada|country_Korea|country_Peru|country_Mongolia|
+---+---------+--------+------+------+------+------+-----------+-----------+-----------+--------------+-------------+------------+----------------+
| 22|        A|  Canada|     1|     0|     0|     0|          1|          0|          0|             1|            0|           0|               0|
| 34|        B|Mongolia|     0|     0|     0|     1|          0|          0|          1|             0|            0|           0|               1|
| 55|        A|    Peru|     0|     0|     1|     0|          1|          0|          0|             0|            0|           1|               0|
| 44|        C|   Korea|     0|     1|     0|     0|          0|          1|          0|             0|            1|           0|               0|
+---+---------+--------+------+------+------+------+-----------+-----------+-----------+--------------+-------------+------------+----------------+

You can then use the same variables_dict to apply the one_hot_encoder to another dataframe.

df2.show()
+---+---------+--------+
|age|education| country|
+---+---------+--------+
| 22|        A|  Canada|
| 66|        B|Mongolia|
| 55|        D|    Peru|
| 44|        C|   China|
+---+---------+--------+

The dataframe above has column variables that have been seen before ie. (66, D, China).

We can use the following code to transform 'df2' to have identical columns as 'df'.

for col in variables_dict.keys():
    for val in variables_dict[col]:
        df2 = df2.withColumn("{}_{}".format(col, val), functions.when((df2[col] == val), 1).otherwise(0))

+---+---------+--------+------+------+------+------+-----------+-----------+-----------+--------------+-------------+------------+----------------+
|age|education| country|age_22|age_44|age_55|age_34|education_A|education_C|education_B|country_Canada|country_Korea|country_Peru|country_Mongolia|
+---+---------+--------+------+------+------+------+-----------+-----------+-----------+--------------+-------------+------------+----------------+
| 22|        A|  Canada|     1|     0|     0|     0|          1|          0|          0|             1|            0|           0|               0|
| 66|        B|Mongolia|     0|     0|     0|     0|          0|          0|          1|             0|            0|           0|               1|
| 55|        D|    Peru|     0|     0|     1|     0|          0|          0|          0|             0|            0|           1|               0|
| 44|        C|   China|     0|     1|     0|     0|          0|          1|          0|             0|            0|           0|               0|
+---+---------+--------+------+------+------+------+-----------+-----------+-----------+--------------+-------------+------------+----------------+
user_147222970
Related