Groupby和collect_列表根据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进行分组，然后在列名上收集列表。为此，我使用了代码

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()