如何在PySpark中为用户事件添加索引并保持排序顺序？

我有如下DataFrame：

+------------+------------------+--------------------+
|id.         |install_time_first|           timestamp|
+------------+------------------+--------------------+
|           2|        2022-02-02|2022-02-01 10:03:...|
|           3|        2022-02-01|2022-02-01 10:00:...|
|           2|        2022-02-02|                null|
|           3|        2022-02-01|2022-02-03 11:35:...|
|           1|        2022-02-01|                null|
|           2|        2022-02-02|2022-02-02 10:05:...|
|           3|        2022-02-01|2022-02-01 10:05:...|
|           4|        2022-02-02|                null|
|           1|        2022-02-01|2022-02-01 10:05:...|
|           2|        2022-02-02|2022-02-02 10:05:...|
|           4|        2022-02-02|2022-02-03 11:35:...|
|           1|        2022-02-01|                null|
|           1|        2022-02-01|2022-02-01 10:03:...|
|           1|        2022-02-01|2022-02-01 10:05:...|
|           4|        2022-02-02|2022-02-03 11:35:...|
|           2|        2022-02-02|2022-02-02 11:00:...|
|           4|        2022-02-02|2022-02-03 11:35:...|
|           3|        2022-02-01|2022-02-04 11:35:...|
|           1|        2022-02-01|2022-02-01 10:00:...|
+------------+------------------+--------------------+

我希望先按install_time_first对该DataFrame排序，再为每个用户的所有事件添加统一的索引，并保持排序后的顺序。示例输出如下：

+------------+------------------+--------------------+-----+
|id.         |install_time_first|           timestamp|index|
+------------+------------------+--------------------+-----+
|           1|        2022-02-01|                null|    1|
|           1|        2022-02-01|                null|    1|
|           1|        2022-02-01|2022-02-01 10:00:...|    1|
|           1|        2022-02-01|2022-02-01 10:03:...|    1|
|           1|        2022-02-01|2022-02-01 10:05:...|    1|
|           1|        2022-02-01|2022-02-01 10:05:...|    1|
|           3|        2022-02-01|2022-02-01 10:00:...|    2|
|           3|        2022-02-01|2022-02-01 10:05:...|    2|
|           3|        2022-02-01|2022-02-03 11:35:...|    2|
|           3|        2022-02-01|2022-02-04 11:35:...|    2|
|           2|        2022-02-02|                null|    3|
|           2|        2022-02-02|2022-02-01 10:03:...|    3|
|           2|        2022-02-02|2022-02-02 10:05:...|    3|
|           2|        2022-02-02|2022-02-02 10:05:...|    3|
|           2|        2022-02-02|2022-02-02 11:00:...|    3|
|           4|        2022-02-02|                null|    4|
|           4|        2022-02-02|2022-02-03 11:35:...|    4|
|           4|        2022-02-02|2022-02-03 11:35:...|    4|
|           4|        2022-02-02|2022-02-03 11:35:...|    4|
+------------+------------------+--------------------+-----+

我无法做到既添加索引又维持排序效果，请问该如何实现？

可以通过以下两步实现，利用pandas的排序、分组和编码功能：

1. 先按install_time_first排序，同时辅助按id.排序，确保同安装时间内的用户事件聚在一起
2. 对排序后的DataFrame，基于install_time_first和id.的组合分组，生成连续的索引值

具体代码如下：

import pandas as pd

# 假设你的DataFrame名为df
# 第一步：按安装时间排序，同时间内按用户ID排序保证分组连续
df_sorted = df.sort_values(by=['install_time_first', 'id.'])

# 第二步：为每个(安装时间, 用户ID)组合分配唯一连续索引，从1开始计数
df_sorted['index'] = df_sorted.groupby(['install_time_first', 'id.']).ngroup() + 1

# 查看结果
print(df_sorted)

代码说明 #

• sort_values(by=['install_time_first', 'id.'])：先按安装时间升序排列，同安装时间内按用户ID排序，确保同一用户的所有事件连续排列，整体顺序符合需求
• groupby(['install_time_first', 'id.']).ngroup() + 1：ngroup()会为每个分组分配从0开始的连续编号，加1后得到从1开始的索引，完全匹配示例中的输出格式

这样处理后，既能保持按install_time_first排序后的顺序，又能为每个用户的所有事件分配统一的索引值。

内容的提问来源于stack exchange，提问作者nirkov