one-hot encode of multiple string categorical features using Spark DataFrames

Question

My goal is to one-hot encode a list of categorical columns using Spark DataFrames. For example, same like get_dummies() function does in Pandas.

The data set, bureau.csv originally have been taken from a Kaggle competition Home Credit Default Risk. Here is my entry table example, say entryData, where it is filtered where only KEY = 100001.

# primary key
KEY = 'SK_ID_CURR'
data = spark.read.csv("bureau.csv", header=True, inferSchema=True)
# sample data from bureau.csv of 1716428 rows
entryData = data.select(columnList).where(F.col(KEY) == 100001).show()
print(entryData)

+----------+-------------+---------------+---------------+
|SK_ID_CURR|CREDIT_ACTIVE|CREDIT_CURRENCY|    CREDIT_TYPE|
+----------+-------------+---------------+---------------+
|    100001|       Closed|     currency 1|Consumer credit|
|    100001|       Closed|     currency 1|Consumer credit|
|    100001|       Closed|     currency 1|Consumer credit|
|    100001|       Closed|     currency 1|Consumer credit|
|    100001|       Active|     currency 1|Consumer credit|
|    100001|       Active|     currency 1|Consumer credit|
|    100001|       Active|     currency 1|Consumer credit|
+----------+-------------+---------------+---------------+

I'm looking to one-hot encode the list columnList by creating the function catg_encode(entryData, columnList),

columnList = cols_type(entryData, obj=True)[1:]
print(columnList)
['CREDIT_ACTIVE', 'CREDIT_CURRENCY', 'CREDIT_TYPE']

Note cols_type() is a function that returns a list of columns either categorical columns (if obj=True) or numerical ones (if obj=False).

I have succeeded in one-hot encode the first column 'CREDIT_ACTIVE' but I couldn't for the hole columns simultaneously, I mean to build the function catg_encode.

# import necessary modules
from pyspark.sql import functions as F

# look for all distinct categoris within a given feature (here 'CREDIT_ACTIVE')
categories = entryData.select(columnList[0]).distinct().rdd.flatMap(lambda x: x).collect()
# one-hot encode the categories
exprs = [F.when(F.col(columnList[0]) == category, 1).otherwise(0).alias(category) for category in categories]
# nice table with encoded feature 'CREDIT_ACTIVE'
oneHotEncode = entryData.select(KEY, *exprs)
print(oneHotEncode)

+----------+--------+----+------+------+
|SK_ID_CURR|Bad debt|Sold|Active|Closed|
+----------+--------+----+------+------+
|    100001|       0|   0|     0|     1|
|    100001|       0|   0|     0|     1|
|    100001|       0|   0|     0|     1|
|    100001|       0|   0|     0|     1|
|    100001|       0|   0|     1|     0|
|    100001|       0|   0|     1|     0|
|    100001|       0|   0|     1|     0|
+----------+--------+----+------+------+

Here the feature 'CREDIT_ACTIVE' has 4 distinct categories; ['Bad debt', 'Sold', 'Active', 'Closed'].

Note I have even tried IndexToString and OneHotEncoderEstimator but didn't help for this specific task.

I am expecting to have the following output,

+----------+--------+----+------+------+----------+----------+----------+----------+----------+---
|SK_ID_CURR|Bad debt|Sold|Active|Closed|currency 1|currency 2|currency 3|currency 4|..........|...
+----------+--------+----+------+------+----------+----------+----------+----------+----------+---
|    100001|       0|   0|     0|     1|         1|         0|         0|         0|        ..|   
|    100001|       0|   0|     0|     1|         1|         0|         0|         0|        ..|
|    100001|       0|   0|     0|     1|         1|         0|         0|         0|        ..|
|    100001|       0|   0|     0|     1|         1|         0|         0|         0|        ..|
|    100001|       0|   0|     1|     0|         1|         0|         0|         0|        ..|
|    100001|       0|   0|     1|     0|         1|         0|         0|         0|        ..|
|    100001|       0|   0|     1|     0|         1|         0|         0|         0|        ..|
+----------+--------+----+------+------+----------+----------+----------+----------+----------+---

The continuous points ... are for the remaining categories of the feature 'CREDIT_TYPE' which are

['Loan for the purchase of equipment', 'Cash loan (non-earmarked)', 'Microloan', 'Consumer credit', 'Mobile operator loan', 'Another type of loan', 'Mortgage', 'Interbank credit', 'Loan for working capital replenishment', 'Car loan', 'Real estate loan', 'Unknown type of loan', 'Loan for business development', 'Credit card', 'Loan for purchase of shares (margin lending)'].

Remarque: I have seen this post E-num / get Dummies in pyspark but doesn't automate the process for many columns, case of big data. The post gives a solution to write separate code for each categorical feature which is not my case problem.

if the order of these distinct categories over columns is not a concern, you can create an array column including all targeted columns i.e. `df.withColumn('arr', F.array(columnList))` and then use **CountVectorizer** to create the one hot encoding in one shot. one of my old posts as an example: https://stackoverflow.com/questions/58010126 — jxc, Nov 22 '19 at 21:33

Napoleon Borntoparty · Answer 1 · 2019-11-22T15:51:55.637

There's two ways to juice this particular lemon. Let's have a look at them.

Pivoting and joining

import pyspark.sql.functions as f

df1 = spark._sc.parallelize([
    [100001, 'Closed', 'currency 1', 'Consumer credit'],
    [100001, 'Closed', 'currency 1', 'Consumer credit'],
    [100001, 'Closed', 'currency 1', 'Consumer credit'],
    [100001, 'Closed', 'currency 1', 'Consumer credit'],
    [100001, 'Active', 'currency 1', 'Consumer credit'],
    [100001, 'Active', 'currency 1', 'Consumer credit'],
    [100001, 'Active', 'currency 1', 'Consumer credit'],
    [100002, 'Active', 'currency 2', 'Consumer credit'],
]).toDF(['SK_ID_CURR', 'CREDIT_ACTIVE', 'CREDIT_CURRENCY', 'CREDIT_TYPE'])

# this can be done dynamically, but I don't have all categories
categories = ['Active', 'Closed', 'Bad debt', 'Sold']

# we need to pivot without aggregation, so I need to add an `id` column and group by it as well
credit_groups = (
  df1.withColumn('id', f.monotonically_increasing_id())
     .groupBy('SK_ID_CURR', 'id')
     .pivot('CREDIT_ACTIVE', values=categories)
     .agg(f.lit(1))
     .drop('id')
)

# currency groups are just a 1 for each currency and ID, as per the example data
# if this is not the case, something more clever needs to be here 
currency_groups = df1.groupBy('SK_ID_CURR').pivot('CREDIT_CURRENCY').agg(f.lit(1))

# join the two pivoted tables on the ID and fill nulls to zeroes
credit_groups.join(currency_groups, on=['SK_ID_CURR'], how='inner').na.fill(0).show()

+----------+------+------+--------+----+----------+----------+
|SK_ID_CURR|Active|Closed|Bad debt|Sold|currency 1|currency 2|
+----------+------+------+--------+----+----------+----------+
|    100002|     1|     0|       0|   0|         0|         1|
|    100001|     0|     1|       0|   0|         1|         0|
|    100001|     1|     0|       0|   0|         1|         0|
|    100001|     1|     0|       0|   0|         1|         0|
|    100001|     0|     1|       0|   0|         1|         0|
|    100001|     0|     1|       0|   0|         1|         0|
|    100001|     1|     0|       0|   0|         1|         0|
|    100001|     0|     1|       0|   0|         1|         0|
+----------+------+------+--------+----+----------+----------+

Using StringIndexer and OneHotEncoderEstimator such as following:

from pyspark.ml import Pipeline
from pyspark.ml.feature import OneHotEncoderEstimator, StringIndexer

indexers = [StringIndexer(inputCol=column, outputCol=column+"_NUMERIC").fit(df1) for column in ['CREDIT_ACTIVE', 'CREDIT_CURRENCY']]

pipeline = Pipeline(stages=indexers)
df_indexed = pipeline.fit(df1).transform(df1)
df_indexed.show()

+----------+-------------+---------------+---------------+---------------------+-----------------------+
|SK_ID_CURR|CREDIT_ACTIVE|CREDIT_CURRENCY|    CREDIT_TYPE|CREDIT_ACTIVE_NUMERIC|CREDIT_CURRENCY_NUMERIC|
+----------+-------------+---------------+---------------+---------------------+-----------------------+
|    100001|       Closed|     currency 1|Consumer credit|                  0.0|                    0.0|
|    100001|       Closed|     currency 1|Consumer credit|                  0.0|                    0.0|
|    100001|       Closed|     currency 1|Consumer credit|                  0.0|                    0.0|
|    100001|       Closed|     currency 1|Consumer credit|                  0.0|                    0.0|
|    100001|       Active|     currency 1|Consumer credit|                  1.0|                    0.0|
|    100001|       Active|     currency 1|Consumer credit|                  1.0|                    0.0|
|    100001|       Active|     currency 1|Consumer credit|                  1.0|                    0.0|
|    100002|       Active|     currency 2|Consumer credit|                  1.0|                    1.0|
+----------+-------------+---------------+---------------+---------------------+-----------------------+

And from here on out, you use one-hot encoding on your newly created numerical columns. I personally recommend route 1, as it is more readable. Route 2, however, allows you to chain OneHotEncoderEstimator into the declared Pipeline as well, making the code executable from one line after declaration. Hope this helps.

@napoleon-borntoparty, thanks for your time and your answer. I can say that with these two approaches we still need to run your code as many times as the number of categorical features we have in the dataset which doesn't solve the problem I asked. — Joe, Nov 22 '19 at 16:48
@Joe ah I see, you want to daisy-chain your categoricals. Yeah that's fair, in which case you can adapt approach 2. and `Pipeline` multiple `StringIndexer`s, followed by `OneHotEncoderEstimator` for as many categoricals as you would have in `for column in ['CREDIT_ACTIVE', 'CREDIT_CURRENCY', '...']` — Napoleon Borntoparty, Nov 22 '19 at 16:55
Thanks, @napoleon-borntoparty. I understand but it doesn't one-hot encode the columns. I mean, it doesn't output the same what I'm expecting (see my post). — Joe, Nov 22 '19 at 17:01

score 0 · Answer 2 · answered Nov 22 '19 at 15:30

The OHE defined in SparkML can only handle one column at a time which can be sub optimal. You can implement this multi column OHE yourself. You were actually on the right track.

import pyspark.sql.functions as F

# let's define some data
l = [('a', 1), ('b', 2), ('c', 1), ('a', 1)]
df = spark.createDataFrame(l, ['c1', 'c2'])
# the list of column we want to encode
cols = ['c1', 'c2']

# defining a struct that associates each column name to its value
col_struct = [
  F.struct(F.lit(c).alias('key'),
           F.col(c).cast('string').alias('value')) for c in cols
]

# Then we explode these struct, group by column name and collect the
# distinct values. Finally, we collect everything to the driver.
ohe_rows = df.distinct()\
  .select(*cols).select(F.explode(F.array(*col_struct)).alias("x"))\
  .groupBy("x.key")\
  .agg(F.collect_set(F.col("x.value")).alias("values"))\
  .collect()

# then we build one spark column per column and per value of that column
# so as to encode the values
ohe = [
          [
              F.when(F.col(row['key']) == value, 1)
               .otherwise(0)
               .alias(row['key']+'_'+value) for value in row['values']
          ] for row in ohe_rows
      ]

# ohe is a list of lists so we use itertools to flatten it
import itertools
ohe_list = list(itertools.chain(*ohe))

# and voila
df.select(* [df.c1, df.c2] + ohe_list).show()
+---+---+----+----+----+----+----+
| c1| c2|c1_c|c1_b|c1_a|c2_1|c2_2|
+---+---+----+----+----+----+----+
|  a|  1|   0|   0|   1|   1|   0|
|  b|  2|   0|   1|   0|   0|   1|
|  c|  1|   1|   0|   0|   1|   0|
|  a|  1|   0|   0|   1|   1|   0|
+---+---+----+----+----+----+----+
# or simply df.select(*ohe_list)

You can chain `OneHotEncoderEstimator`s into a `Pipeline`, similarly to what I've done with `StringIndexer` below. — Napoleon Borntoparty, Nov 22 '19 at 15:33
Right, but the question was about making the computations in parallel and I don't think that Spark makes that optimization when you chain the transformations. — Oli, Nov 23 '19 at 10:58

one-hot encode of multiple string categorical features using Spark DataFrames

2 Answers2