Creating a dictionary type column in dataframe

Question

Consider the following dataframe:

------------+--------------------+
|id|          values
+------------+--------------------+
|          39|a,a,b,b,c,c,c,c,d
|         520|a,b,c
|         832|a,a

I want to convert it into the following DataFrame:

------------+--------------------+
|id|          values
+------------+--------------------+
|          39|{"a":2, "b": 2,"c": 4,"d": 1}
|         520|{"a": 1,"b": 1,"c": 1}
|         832|{"a": 2}

I tried two approaches:

1. Converting the dataframe to rdd. Then I mapped the value column to a frequancy counter function. But I get errors on converting the rdd back to the dataframe

2. Using a udf to essentially do the same thing as above.

The reason I want to have a dictionary column is to load it as a json in one of my python application.

Answer 1

You can do this with a udf that returns a MapType column.

from pyspark.sql.types import MapType, StringType, IntegerType
from collections import Counter

my_udf = udf(lambda s: dict(Counter(s.split(','))), MapType(StringType(), IntegerType()))
df = df.withColumn('values', my_udf('values'))
df.collect()

[Row(id=39, values={u'a': 2, u'c': 4, u'b': 2, u'd': 1}),
 Row(id=520, values={u'a': 1, u'c': 1, u'b': 1}),
 Row(id=832, values={u'a': 2})]