Groupby и collect_list поддерживают порядок на основе другого столбца в PySpark

Question

У меня есть такой фрейм данных PySpark,

+----------+------------+------------+------------+
|   Name   |  dateCol1  |  dateCol2  |  dateCol3  |
+----------+------------+------------+------------+
| user1    | 2018-01-01 | 2018-01-10 | 2018-01-01 |
| user1    | 2018-01-11 | 2018-01-20 | 2018-01-01 |
| user2    | 2018-01-11 | 2018-01-20 | 2018-01-11 |
| user1    | 2019-01-21 | 2018-01-30 | 2018-01-01 |
+----------+------------+------------+------------+

Я хочу сгруппировать этот набор данных по ключам, dateCol1 и dateCol2 и, таким образом, collect_list по столбцу Name. Для этого я использую код:

spark_df.groupBy('dateCol1', 'dateCol2').agg(F.collect_list('Name'))

При сборе столбца в список я также хочу поддерживать порядок значений на основе столбца dateCol3.

Например, Я хочу убедиться, что для dateCol1 == '2018-01-11' и dateCol2 == '2018-01-20', собирающих в список, я всегда получу [user1, user2] (на основе порядка dateCol3).

Требуемый вывод для фрейма данных:

+------------+------------+----------------+
|  dateCol1  |  dateCol2  |    List        |
+------------+------------+----------------+
| 2018-01-01 | 2018-01-10 | [user1]        |
| 2018-01-11 | 2018-01-20 | [user1, user2] |
| 2019-01-21 | 2018-01-30 | [user1]        |
+------------+------------+----------------+

collect_list не будет поддерживать порядок по умолчанию. Как убедиться, что собранный список упорядочен на основе другого внешнего столбца из фрейма данных?

Answer 1

Вы можете попробовать:

spark_df.orderBy('dateCol3', ascending=True).groupBy('dateCol1', 'dateCol2').agg(F.collect_list('Name'))

В качестве альтернативы, хотя это будет немного излишним, вы можете использовать оконное управление:

from pyspark.sql import Window as w

spark_df.select('dateCol1', 'dateCol2', F.collect_list('Name').over(w.partitionBy(['dateCol1','dateCol2']).orderBy(F.col('dateCol3'))).alias('Name')).distinct()