Pyspark collect_set или collect_list с groupby

Ответ 1

Вам нужно использовать агг. Пример:

from pyspark import SparkContext
from pyspark.sql import HiveContext
from pyspark.sql import functions as F

sc = SparkContext("local")

sqlContext = HiveContext(sc)

df = sqlContext.createDataFrame([
    ("a", None, None),
    ("a", "code1", None),
    ("a", "code2", "name2"),
], ["id", "code", "name"])

df.show()

+---+-----+-----+
| id| code| name|
+---+-----+-----+
|  a| null| null|
|  a|code1| null|
|  a|code2|name2|
+---+-----+-----+

Обратите внимание, что выше необходимо создать HiveContext. См. fooobar.com/questions/358663/... о работе с различными версиями Spark.

(df
  .groupby("id")
  .agg(F.collect_set("code"),
       F.collect_list("name"))
  .show())

+---+-----------------+------------------+
| id|collect_set(code)|collect_list(name)|
+---+-----------------+------------------+
|  a|   [code1, code2]|           [name2]|
+---+-----------------+------------------+

Ответ 2

Если ваш фрейм данных большой, вы можете попробовать использовать pandas udf (GROUPED_AGG), чтобы избежать ошибки памяти. Это также намного быстрее.

Сгруппированные UDF-функции Pandas аналогичны агрегатным функциям Spark. Сгруппированные пользовательские функции Pandas используются с groupBy(). Agg() и pyspark.sql.Window. Он определяет агрегацию от одного или нескольких pandas.Series до скалярного значения, где каждый pandas.Series представляет столбец в группе или окне. pandas udf

Пример:

import pyspark.sql.functions as F

@F.pandas_udf('string', F.PandasUDFType.GROUPED_AGG)
def collect_list(name):
    return ', '.join(name)

grouped_df = df.groupby('id').agg(collect_list(df["name"]).alias('names'))