5

enter image description here

I am hoping to dummy encode my categorical variables to numerical variables like shown in the image below, using Pyspark syntax.

I read in data like this

data = sqlContext.read.csv("data.txt", sep = ";", header = "true")

In python I am able to encode my variables using the below code

data = pd.get_dummies(data, columns = ['Continent'])

However I am not sure how to do it in Pyspark.

Any assistance would be greatly appreciated.

zero323
  • 322,348
  • 103
  • 959
  • 935
ALK
  • 87
  • 1
  • 2
  • 9

1 Answers1

11

Try this:

import pyspark.sql.functions as F 
categ = df.select('Continent').distinct().rdd.flatMap(lambda x:x).collect()
exprs = [F.when(F.col('Continent') == cat,1).otherwise(0)\
            .alias(str(cat)) for cat in categ]
df = df.select(exprs+df.columns)

Exclude df.columns if you do not want the original columns in your transformed dataframe.

mayank agrawal
  • 2,495
  • 2
  • 13
  • 32