求带过滤条件的SQL COUNT窗口函数的Pandas等价实现方案

需求对应SQL逻辑 #

目标是实现以下SQL的功能：

select t.*,
    count(*) filter(where grp = 'new') over(partition by usr order by id) rn
from mytable t
order by usr, id

这段SQL会为每一行计算截至当前行（按usr分组、id排序），grp等于'new'的累计行数。

原代码问题 #

你之前尝试的代码：

mytable['rn'] = mytable.groupby('usr')['grp'].transform('count')

仅统计了每个usr分组的总记录数，既没有过滤grp='new'的条件，也没有实现累计计数的窗口函数逻辑，因此不符合需求。

正确的向量化实现方案 #

方法1：分步实现（清晰直观）

先确保数据已按usr和id排序（和SQL逻辑对齐），再通过标记+分组累计求和实现：

# 1. 先按usr和id排序，保证累计顺序正确
mytable = mytable.sort_values(['usr', 'id'])
# 2. 标记grp为'new'的行
mytable['is_new'] = mytable['grp'] == 'new'
# 3. 按usr分组，对标记列做累计求和，得到目标rn列
mytable['rn'] = mytable.groupby('usr')['is_new'].cumsum()
# 可选：删除临时标记列
# mytable.drop('is_new', axis=1, inplace=True)

方法2：一步到位（更简洁）

无需临时列，直接通过transform结合cumsum实现：

# 先排序（若未排序）
mytable = mytable.sort_values(['usr', 'id'])
# 直接计算累计计数
mytable['rn'] = mytable.groupby('usr')['grp'].transform(lambda s: (s == 'new').cumsum())

逻辑说明 #

将grp='new'的行转为数值1，其余转为0，在每个usr分组内按id的顺序累计求和，完全等价于SQL中带过滤条件的COUNT窗口函数逻辑，且全程是向量化操作，效率较高。

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