Convert Pyspark dataframe to dictionary

Question

I'm trying to convert a Pyspark dataframe into a dictionary.

Here's the sample CSV file -

Col0, Col1
-----------
A153534,BDBM40705
R440060,BDBM31728
P440245,BDBM50445050

I've come up with this code -

from rdkit import Chem
from pyspark import SparkContext
from pyspark.conf import SparkConf
from pyspark.sql import SparkSession

sc = SparkContext.getOrCreate()
spark = SparkSession(sc)

df = spark.read.csv("gs://my-bucket/my_file.csv") # has two columns

# Creating list
to_list = map(lambda row: row.asDict(), df.collect())

#Creating dictionary
to_dict = {x['col0']: x for x in to_list }

This creates a dictionary like below -

'A153534': {'col0': 'A153534', 'col1': 'BDBM40705'}, 'R440060': {'col0': 'R440060', 'col1': 'BDBM31728'}, 'P440245': {'col0': 'P440245', 'col1': 'BDBM50445050'} 

But I want a dictionary like this -

{'A153534': 'BDBM40705'}, {'R440060': 'BDBM31728'}, {'P440245': 'BDBM50445050'}

How can I do that?

I tried the rdd solution by Yolo but I'm getting error. Can you please tell me what I am doing wrong?

py4j.protocol.Py4JError: An error occurred while calling o80.isBarrier. Trace: py4j.Py4JException: Method isBarrier([]) does not exist at py4j.reflection.ReflectionEngine.getMethod(ReflectionEngine.java:318) at py4j.reflection.ReflectionEngine.getMethod(ReflectionEngine.java:326) at py4j.Gateway.invoke(Gateway.java:274) at py4j.commands.AbstractCommand.invokeMethod(AbstractCommand.java:132) at py4j.commands.CallCommand.execute(CallCommand.java:79) at py4j.GatewayConnection.run(GatewayConnection.java:238) at java.lang.Thread.run(Thread.java:748)

Answer 1

Here's a way of doing it using rdd:

df.rdd.map(lambda x: {x.Col0: x.Col1}).collect()

[{'A153534': 'BDBM40705'}, {'R440060': 'BDBM31728'}, {'P440245': 'BDBM50445050'}]

Answer 2

This could help you:

df = spark.read.csv('/FileStore/tables/Create_dict.txt',header=True)

df = df.withColumn('dict',to_json(create_map(df.Col0,df.Col1)))

df_list = [row['dict'] for row in df.select('dict').collect()]

df_list

Output is:

['{"A153534":"BDBM40705"}',
 '{"R440060":"BDBM31728"}',
 '{"P440245":"BDBM50445050"}']