I am using Ipython notebook to work with pyspark applications. I have a CSV file with lots of categorical columns to determine whether the income falls under or over the 50k range. I would like to perform a classification algorithm taking all the inputs to determine the income range. I need to build a dictionary of variables to mapped variables and use a map function to map the variables to numbers for processing. Essentially, I would my dataset to be in a numerical format so that I can work on implementing the models.

In the data set, there are categorical columns like education, marital status, working class etc. Can someone tell me how to convert them into numerical columns in pyspark?

workclass = {'?':0,'Federal-gov':1,'Local-gov':2,'Never-  worked':3,'Private':4,'Self-emp-inc':5,'Self-emp-not-inc':6,'State-gov':7,'Without-pay':8}

I created a sample dictionary with key value pairs for work class. But, I don't know how to use this in a map function and replace the categorical data in the CSV file with the corresponding value.

wc = pd.read_csv('PATH', usecols = ['Workclass'])

df = pd.DataFrame(wc)
wcdict = {' ?':0,' Federal-gov':1,' Local-gov':2,' Never-worked':3,' Private':4,' Self-emp-inc':5,' Self-emp-n-inc':6,' State-gov':7,' Without-pay':8}
df_new = df.applymap(lambda s: wcdict.get(s) if s in wcdict else s)

This is the code I have written in normal python to convert the categorical data into numerical data. It works fine. I want to do the conversion in spark context. And, there are 9 categorical columns in the data source. Is there a way to automate the dictionary update process to have a KV pair for all 9 columns?


3 Answers 3


This can be done using StringIndexer in PySpark and the reverse using IndexToString for reference please check this:

from pyspark.ml.feature import StringIndexer

df = sqlContext.createDataFrame(
    [(0, "a"), (1, "b"), (2, "c"), (3, "a"), (4, "a"), (5, "c")],
    ["id", "category"])
indexer = StringIndexer(inputCol="category", outputCol="categoryIndex")
indexed = indexer.fit(df).transform(df)

For more details, please check the spark documentation

workclass = {'?':0,'Federal-gov':1,'Local-gov':2,'Never-  worked':3,'Private':4,'Self-emp-inc':5,'Self-emp-not-inc':6,'State-gov':7,'Without-pay':8}

try defining a mapper fuction which return key :

def mapr(dict_key):
    return workclass[dict_key]

print list(map(mapr,workclass))
  • $\begingroup$ Hey, could you please explain to me what this block does? I ran my script with this code added and I got [6, 1, 4, 3, 5, 7, 8, 0, 2] as the output. I want to substitute numerical values to the work class content using the values in the dictionary. $\endgroup$
    – SRS
    Jun 30, 2015 at 15:43
  • $\begingroup$ Hi, The mapr function will return numerical value associated with the category value. eg : 6 for 'Self-emp-not-inc', python dictionaries are unordered. If you want an ordered dictionary, try collections.OrderedDict. $\endgroup$ Jun 30, 2015 at 16:35
  • $\begingroup$ Okay, now I understand the function. The thing is, I have a CSV with several thousand rows and there is a column named Workclass which contains any one of the value mentioned in the dictionary. So, for each row, I need to change the text in that column to a number by comparing the text with the dictionary and substitute the corresponding number. How do I use a function to parse the column by rows and compare the values with the dictionary? $\endgroup$
    – SRS
    Jun 30, 2015 at 21:16
  • $\begingroup$ You can create an additional column, say 'workclass_num' which store numerical values corresponding to the categorical value. Check Python Pandas library. $\endgroup$ Jul 5, 2015 at 8:32

This could work. Suppose if your column name is "Marital status" and categorical,

v1=dataset['Marital status'].unique()

dataset['Marital status'].replace(to_replace=v1,value= list(range(len(v1))), inplace=True)

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.