Как группировать по одному столбцу и брать минимум по другим столбцам?

Question

Input DataFrame

+----+---+----+--------+
|type|aud|code|priority|
+----+---+----+--------+
|   A| AA| qwe|       2|
|   B| AA| qwe|       1|
|   A| BB| qwe|       2|
|   B| BB| qwe|       1|
|   A| AA| jku|       2|
|   B| AA| jku|       1|
|   A| BB| jku|       2|
|   B| BB| jku|       1|
+----+---+----+--------+

Output DataFrame

+----+---+----+--------+
|type|aud|code|priority|
+----+---+----+--------+
|   B| AA| qwe|       1|
|   B| AA| jku|       1|
|   B| BB| qwe|       1|
|   B| BB| jku|       1|
+----+---+----+--------+

Шаги описания проблемы:

1. Сначала выполните groupBy в столбце "aud"

2. Затем в каждой группе для каждого кода ""value найти свое минимальное значение в столбце" priority "

3. Предположим, что в группе AA минимальное значение кода" qwe ", использующего столбец" priority ", равно 1, поэтому извлеките его и отфильтруйте другие строки,

Я новичок в pyspark, так что любой может рассказать, как это сделать, используя оконные функции или любым другим оптимизированным способом.Спасибо за чтение.

Answer 1

Если вам нужно сгруппировать по «aud», а затем найти минимальный «приоритет» для «code», вы можете сгруппировать по столбцам «aud» и «code» и найти минимальный «приоритет».

window_spec = Window \
    .partitionBy("aud", "code") \
    .orderBy("priority")

res = df \
    .withColumn("row_num", row_number().over(window_spec)) \
    .where(col("row_num") == 1) \
    .drop("row_num")

res.show()

Выход:

+----+---+----+--------+
|type|aud|code|priority|
+----+---+----+--------+
|   B| AA| jku|       1|
|   B| AA| qwe|       1|
|   B| BB| qwe|       1|
|   B| BB| jku|       1|
+----+---+----+--------+