如何在 pyspark 中創建具有兩個 dataframe 列的字典？

Question

我有一個 dataframe 有兩列，如下所示：

    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

我需要為 Group 列中的每個唯一值遍歷這些數據並執行一些處理。 我正在考慮創建一個字典，其中每個組名稱作為鍵，它們對應的主題列表作為值。

所以，我預期的 output 如下所示：

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

如何在 pyspark 中實現這一點？

Answer 1

看看這個：你可以做groupBy並使用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

如果您需要獨特的主題，則可以使用 collect_set，否則使用 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"))

結果：

+-----+------------------------+-------------------------------+
|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]]|
+-----+------------------------+-------------------------------+