PySpark df to dict: one column as key, the other as value

Viewed 1044

I have a PySpark Data-frame like this:

+----+-----------------+
|  id|    director_name|
+----+-----------------+
| 123|    james cameron|
|  32|   gore verbinski|
|  34|       sam mendes|
|2345|christopher nolan|
| 987|      doug walker|
+----+-----------------+

I need a dict like this:

{123: james cameron, 32: gore verbinski, ...}

I have found a solution using pandas df, but I need to work as parallelized as possible, therefore I'm looking for a spark solution...

2 Answers

RDDs have built in function asDict() that allows to represent each row as a dict.

A Spark Dataframe sparkDF, can converted to an rdd & apply asDict() to get the desired result

Another way can be to convert your Spark DataFrame toPandas and apply as demonstrated under the link - here

Data Preparation

input_str = """
| 123|    james cameron
|  32|   gore verbinski
|  34|       sam mendes
|2345|christopher nolan
| 987|      doug walker
""".split("|")


input_values = list(map(lambda x: x.strip() if x.strip() != 'null' else None, input_str[1:]))

n = len(input_values)

cols = ['id','director_name']

input_list = [tuple(input_values[i:i+2]) for i in range(0,n,2)]

sparkDF = sql.createDataFrame(input_list, cols)

sparkDF.show()

+----+-----------------+
|  id|    director_name|
+----+-----------------+
| 123|    james cameron|
|  32|   gore verbinski|
|  34|       sam mendes|
|2345|christopher nolan|
| 987|      doug walker|
+----+-----------------+

asDict

imm_dict = sparkDF.rdd.map(lambda row: row.asDict()).collect()

id_dict = {d['id']: d['director_name'] for d in imm_dict}

id_dict

{'123': 'james cameron',
 '32': 'gore verbinski',
 '34': 'sam mendes',
 '2345': 'christopher nolan',
 '987': 'doug walker'}

ToPandas

id_dict = sparkDF.toPandas().set_index('id').T.to_dict('records')[0]

{'123': 'james cameron',
 '32': 'gore verbinski',
 '34': 'sam mendes',
 '2345': 'christopher nolan',
 '987': 'doug walker'}

Use spark inbuilt functions map_from_arrays, to_json, collect_list with groupBy.

Example:

sparkDF.\
groupBy(lit(1)).\
agg(to_json(map_from_arrays(collect_list(col("id")),collect_list(col("director_name")))).alias("dict")).\
select("dict").\
collect()[0][0]
#'{"123":"james cameron","32":"gore verbinski","34":"sam mendes","2345":"christopher nolan","987":"doug walker"}'
Related