如何在Pandas中实现类似SQL的ROW_NUMBER()分组排序功能

数据集 #

timestamp   conversationId   UserId  MessageId       tpMessage   Message 
1614578324  ceb9004ae9d3    1c376ef 5bbd34859329    question    Where do you live?
1614578881  ceb9004ae9d3    1c376ef d3b5d3884152    answer      Brooklyn
1614583764  ceb9004ae9d3    1c376ef 0e4501fcd61f    question    What's your name?
1614590885  ceb9004ae9d3    1c376ef 97d841b79ff7    answer      Phill
1614594952  ceb9004ae9d3    1c376ef 11ed3fd24767    question    What's your gender?
1614602036  ceb9004ae9d3    1c376ef 601538860004    answer      Male
1614602581  ceb9004ae9d3    1c376ef 8bc8d9089609    question    How old are you?
1614606219  ceb9004ae9d3    1c376ef a2bd45e64b7c    answer      35
1614606240  jto9034pe0i5    1c489rl o6bd35e64b5j    question    What's your name?
1614606250  jto9034pe0i5    1c489rl 96jd89i55b7t    answer      Robert  

需求 #

实现与以下SQL语句等价的功能：

ROW_NUMBER() OVER(PARTITION BY userId ORDER BY UserId,timestamp,conversationId ASC) AS num_Row

尝试过的错误方法 #

1. 错误将排序字段加入分组维度：

df['row_number'] = df.groupby(['userId','timestamp','conversationId']).cumcount() + 1

2. 排序参数设置错误（timestamp设为降序）：

df['row_number'] = df.sort_values(['userId','timestamp','conversationId'], ascending=[True,False]) \
             .groupby(['userId']) \
             .cumcount() + 1
print(df)

预期输出 #

timestamp   conversationId   UserId  MessageId       tpMessage   Message                num_row     
1614578324  ceb9004ae9d3    1c376ef 5bbd34859329    question    Where do you live?  1
1614578881  ceb9004ae9d3    1c376ef d3b5d3884152    answer      Brooklyn            2
1614583764  ceb9004ae9d3    1c376ef 0e4501fcd61f    question    What's your name?   3
1614590885  ceb9004ae9d3    1c376ef 97d841b79ff7    answer      Phill               4
1614594952  ceb9004ae9d3    1c376ef 11ed3fd24767    question    What's your gender? 5
1614602036  ceb9004ae9d3    1c376ef 601538860004    answer      Male                6
1614602581  ceb9004ae9d3    1c376ef 8bc8d9089609    question    How old are you?    7
1614606219  ceb9004ae9d3    1c376ef a2bd45e64b7c    answer      35                  8
1614606240  jto9034pe0i5    1c489rl o6bd35e64b5j    question    What's your name?   1
1614606250  jto9034pe0i5    1c489rl 96jd89i55b7t    answer      Robert              2

解决方案 #

错误原因 #

• 第一种方法：分组时错误包含了timestamp和conversationId，SQL中仅按userId分区，这两个字段是排序依据而非分组依据。
• 第二种方法：排序时ascending参数设置错误，需求是升序（ASC），但代码中把timestamp设为降序，导致行号顺序不符合预期。

正确代码 #

# 先按指定字段升序排序
df_sorted = df.sort_values(['userId', 'timestamp', 'conversationId'], ascending=True)
# 按userId分组后生成行号
df_sorted['num_row'] = df_sorted.groupby('userId').cumcount() + 1
# 若需要恢复原数据的顺序，按原索引排序
df = df_sorted.sort_index()

或者更简洁的链式写法：

df['num_row'] = df.sort_values(['userId', 'timestamp', 'conversationId']) \
                  .groupby('userId') \
                  .cumcount() + 1
# 恢复原顺序（可选）
df = df.sort_index()

执行后即可得到与预期一致的结果。

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