Объединить столбцы данных, состоящие из нескольких значений - Spark

Question

У меня есть два фрейма данных Spark с одинаковым столбцом идентификаторов:

df1:

+------+---------+---------+
|ID    |  Name1  |  Name2  |
+------+---------+---------+
|   1  |     A   |     B   |
|   2  |     C   |     D   |
|   3  |     E   |     F   |
+------+---------+---------+

df2:

+------+-------+
|ID    |  key  |
+------+-------+
|   1  |   w   |
|   1  |   x   |
|   2  |   y   |
|   3  |   z   |
+------+-------+

Теперь яхотите создать новый столбец в df1, содержащий все ключевые значения, обозначенные в df2.Итак, я стремлюсь к результату:

+------+---------+---------+---------+
|ID    |  Name1  |  Name2  |  keys   |
+------+---------+---------+---------+
|   1  |     A   |     B   |   w,x   |
|   2  |     C   |     D   |    y    |
|   3  |     E   |     F   |    z    |
+------+---------+---------+---------+

В конечном счете, я хочу найти решение для произвольного количества ключей.

Моя попытка в PySpark:

def get_keys(id):
    x = df2.where(df2.ID == id).select('key')

    return x

df_keys = df1.withColumn("keys", get_keys(col('ID')))

В приведенном выше коде x - это фрейм данных.Поскольку вторым аргументом функции .withColumn должна быть переменная типа Column, я не уверен, как правильно изменить x.

Answer 1

Вы ищете функцию collect_list.

from pyspark.sql.functions import collect_list

df3 = df1.join(df2, df1.ID == df2.ID).drop(df2.ID)

df3.groupBy('ID','Name1','Name2').agg(collect_list('key').alias('keys')).show()

#+---+-----+-----+------+
#| ID|Name1|Name2|  keys|
#+---+-----+-----+------+
#|  1|    A|    B|[w, x]|
#|  3|    C|    F|   [z]|
#|  2|    B|    D|   [y]|
#+---+-----+-----+------+

Если вам нужны только уникальные ключи, вы можете использовать collect_set