Агрегирование и создание массива объекта Dictionary в столбце Spark Dataframe

Question

Я создал фрейм данных игрушечной искры:

import numpy as np
import pyspark
from pyspark.sql import functions as sf
from pyspark.sql import functions as F

# sc = pyspark.SparkContext()
# sqlc = pyspark.SQLContext(sc)
df = spark.createDataFrame([('csc123','sr1', 'tac1', 'abc'), 
                            ('csc123','sr2', 'tac1', 'abc'), 
                            ('csc234','sr3', 'tac2', 'bvd'),
                            ('csc345','sr5', 'tac2', 'bvd')
                           ], 
                           ['bug_id', 'sr_link', 'TAC_engineer','de_manager'])
df.show()

+------+-------+------------+----------+
|bug_id|sr_link|TAC_engineer|de_manager|
+------+-------+------------+----------+
|csc123|    sr1|        tac1|       abc|
|csc123|    sr2|        tac1|       abc|
|csc234|    sr3|        tac2|       bvd|
|csc345|    sr5|        tac2|       bvd|
+------+-------+------------+----------+

Затем я попытался агрегировать и сгенерировать массив [sr_link, sr_link] для каждого идентификатора ошибки

#df = spark.createDataFrame([('row11','row12'), ('row21','row22')], ['colname1', 'colname2'])

df_drop_dup = df.select('bug_id', 'de_manager').dropDuplicates()

df = df.withColumn('joined_column', 
                    sf.concat(sf.col('sr_link'),sf.lit(' '), sf.col('TAC_engineer')))

df_sev_arr = df.groupby("bug_id").agg(F.collect_set("joined_column")).withColumnRenamed("collect_set(joined_column)","sr_array")

df = df_drop_dup.join(df_sev_arr, on=['bug_id'], how='inner')

df.show()

Вот результат :

+------+----------+--------------------+
|bug_id|de_manager|            sr_array|
+------+----------+--------------------+
|csc345|       bvd|          [sr5 tac2]|
|csc123|       abc|[sr2 tac1, sr1 tac1]|
|csc234|       bvd|          [sr3 tac2]|
+------+----------+--------------------+

Но фактический результат, которого я действительно ожидаю, будет примерно таким:

+------+----------+----------------------------------------------------------------------+
|bug_id|de_manager|                                                              sr_array|
+------+----------+----------------------------------------------------------------------+
|csc345|       bvd|                                   [{sr_link: sr5, TAC_engineer:tac2}]|
|csc123|       abc|[{sr_link: sr2, TAC_engineer:tac1},{sr_link: sr1, TAC_engineer: tac1}]|
|csc234|       bvd|                                  [{sr_link: sr3, TAC_engineer: tac2}]|
+------+----------+----------------------------------------------------------------------+

, поскольку я хочу, чтобы окончательный результат можно было сохранить в формате JSON, например:

'bug_id': 'csc123'
'de_manager': 'abc'
'sr_array':
     'sr_link': 'sr2', 'TAC_engineer': 'tac1'
     'sr_link': 'sr1', 'TAC_engineer': 'tac1'

Кто нибудь может помочь? Извините, я очень не знаком с MapType в Spark Dataframe.

Answer 1

Просто изменил несколько функций и добавил новые функции в соответствии с вашими требованиями.

Первая часть останется прежней.

from pyspark.sql import functions as F

# sc = pyspark.SparkContext()
# sqlc = pyspark.SQLContext(sc)
df = spark.createDataFrame([('csc123','sr1', 'tac1', 'abc'), 
                            ('csc123','sr2', 'tac1', 'abc'), 
                            ('csc234','sr3', 'tac2', 'bvd'),
                            ('csc345','sr5', 'tac2', 'bvd')
                           ], 
                           ['bug_id', 'sr_link', 'TAC_engineer','de_manager'])
df.show()

Я только что изменил вторую часть.

>>> df_drop_dup = df.select('bug_id', 'de_manager').dropDuplicates()

Изменена функция переименования с WithcolumnRenamed to Alias и добавлена ​​функция to_json and Struct для получения желаемого результата, а также небольшая модификация фрейма данных naming df -> df1

>>> df1 = df.withColumn('joined_column', F.to_json(F.struct(F.col('sr_link'), F.col('TAC_engineer'))))

>>> df_sev_arr = df1.groupby("bug_id").agg(F.collect_set("joined_column").alias("sr_array"))

>>> df = df_drop_dup.join(df_sev_arr, on=['bug_id'], how='inner')

>>> df.show(truncate=False)
+------+----------+----------------------------------------------------------------------------------+
|bug_id|de_manager|sr_array                                                                          |
+------+----------+----------------------------------------------------------------------------------+
|csc345|bvd       |[{"sr_link":"sr5","TAC_engineer":"tac2"}]                                         |
|csc123|abc       |[{"sr_link":"sr1","TAC_engineer":"tac1"}, {"sr_link":"sr2","TAC_engineer":"tac1"}]|
|csc234|bvd       |[{"sr_link":"sr3","TAC_engineer":"tac2"}]                                         |
+------+----------+----------------------------------------------------------------------------------+

Если у вас есть какие-либо вопросы, связанные с тем же, дайте мне знать.