How to create a dictionary with two dataframe columns in pyspark?

Question

I have a dataframe with two columns that looks as follows:

    df = spark.createDataFrame([('A', 'Science'),
 ('A', 'Math'),
 ('A', 'Physics'),
 ('B', 'Science'),
 ('B', 'English'),
 ('C', 'Math'),
 ('C', 'English'),
 ('C', 'Latin')],
 ['Group', 'Subjects'])


Group   Subjects
A       Science
A       Math
A       Physics
B       Science
B       English
C       Math
C       English
C       Latin

I need to iterate through this data for each unique value in Group column and perform some processing. I'm thinking of creating a dictionary with the each Group name as the key and their corresponding list of Subjects as the value.

So, my expected output would look like below:

{A:['Science', 'Math', 'Physics'], B:['Science', 'English'], C:['Math', 'English', 'Latin']}

How to achieve this in pyspark?

Answer 1

Check this out: You can do groupBy and use collect_list.

    #Input DF
    # +-----+-------+
    # |group|subject|
    # +-----+-------+
    # |    A|   Math|
    # |    A|Physics|
    # |    B|Science|
    # +-----+-------+

    df1 = df.groupBy("group").agg(F.collect_list("subject").alias("subject")).orderBy("group")

    df1.show(truncate=False)

    # +-----+---------------+
    # |group|subject        |
    # +-----+---------------+
    # |A    |[Math, Physics]|
    # |B    |[Science]      |
    # +-----+---------------+

    dict = {row['group']:row['subject'] for row in df1.collect()}

    print(dict)

    # {'A': ['Math', 'Physics'], 'B': ['Science']}

Answer 2

You can use collect_set incase you need unique subjects, else collect_list.

import pyspark.sql.functions as F
 df = spark.createDataFrame([('A', 'Science'),
 ('A', 'Math'),
 ('A', 'Physics'),
 ('B', 'Science'),
 ('B', 'English'),
 ('C', 'Math'),
 ('C', 'English'),
 ('C', 'Latin')],
 ['Group', 'Subjects'])
 
 df_tst=df.groupby('Group').agg(F.collect_set("Subjects").alias('Subjects')).withColumn("dict",F.create_map('Group',"Subjects"))

results:

+-----+------------------------+-------------------------------+
|Group|Subjects                |dict                           |
+-----+------------------------+-------------------------------+
|C    |[Math, Latin, English]  |[C -> [Math, Latin, English]]  |
|B    |[Science, English]      |[B -> [Science, English]]      |
|A    |[Math, Physics, Science]|[A -> [Math, Physics, Science]]|
+-----+------------------------+-------------------------------+