如何在pandas RollingGroupBy中使用size聚合替代count？

我们可以执行如下Pandas操作实现滚动计数：

df.groupby.rolling.agg({'any_df_col': 'count'})

但count生成的是符合分组条件的行的滚动累计序列（如1,1,1,2,3...），而我需要获取每个滚动窗口内对应分组的总行数（即示例中的1,1,3,3,3）。

在Pandas中，通常size可替代count实现这类需求，但RollingGroupby对象不支持size作为聚合函数。以下代码可复现问题：

可正常运行的Count聚合示例 #

import datetime as dt
import pandas as pd

df = pd.DataFrame({'time_ref': [
    dt.datetime(2023, 1, 1, 0, 30),
    dt.datetime(2023, 1, 1, 0, 30),
    dt.datetime(2023, 1, 1, 1),
    dt.datetime(2023, 1, 1, 2),
    dt.datetime(2023, 1, 1, 2, 15),
    dt.datetime(2023, 1, 1, 2, 16),
    dt.datetime(2023, 1, 1, 4),
],
    'value': [1, 2, 1, 10, 10, 10, 10],
    'type': [0, 0, 0, 0, 0, 0, 0]
})
df = df.set_index(pd.DatetimeIndex(df['time_ref']), drop=True)
by = ['value']
window = '1H'
gb_rolling = df.groupby(by=by).rolling(window=window)
agg_d = {'type': 'count'}
test = gb_rolling.agg(agg_d)
print(test)

运行结果：

type
value time_ref                 
1     2023-01-01 00:30:00   1.0
      2023-01-01 01:00:00   2.0
2     2023-01-01 00:30:00   1.0
10    2023-01-01 02:00:00   1.0
      2023-01-01 02:15:00   2.0
      2023-01-01 02:16:00   3.0
      2023-01-01 04:00:00   1.0

报错的Size聚合尝试 #

# 此代码报错
agg_d = {'type': 'size'}
test = gb_rolling.agg(agg_d)
# AttributeError: 'size' is not a valid function for 'RollingGroupby' object

期望输出 #

type
value time_ref
1     2023-01-01 00:30:00   2
      2023-01-01 01:00:00   2
2     2023-01-01 00:30:00   1
10    2023-01-01 02:00:00   3
      2023-01-01 02:15:00   3
      2023-01-01 02:16:00   3
      2023-01-01 04:00:00   1

由于数据窗口由事件自身时间而非日历时间决定，需基于Rolling实现，也接受无需Rolling的替代方案。

解决方案 #

方法1：自定义聚合函数

利用lambda函数在滚动窗口内计算总行数，替代size：

agg_d = {'type': lambda x: len(x)}
test = gb_rolling.agg(agg_d)
test = test.astype(int)  # 转为整数类型，匹配期望输出
print(test)

方法2：使用transform结合滚动窗口

通过transform直接为每个行返回对应滚动窗口的大小：

df['type'] = df.groupby('value').rolling(window)['type'].transform(lambda x: len(x))
test = df.set_index(['value', 'time_ref'])['type'].astype(int)
print(test)

方法3：预计算分组时间范围匹配窗口

如果需要更灵活的窗口逻辑，可以先为每个行计算窗口时间区间，再统计对应分组内落在该区间的行数：

# 按value分组存储时间序列
groups = df.groupby('value')['time_ref']

# 定义函数计算单行对应窗口内的分组行数
def get_window_size(row):
    group = groups.get_group(row['value'])
    start = row['time_ref'] - pd.Timedelta(window)
    end = row['time_ref']
    return len(group[(group >= start) & (group <= end)])

df['type'] = df.apply(get_window_size, axis=1)
test = df.set_index(['value', 'time_ref'])['type']
print(test)

以上方法均可得到期望的输出结果，其中方法1和2更贴合Rolling的原生使用场景，方法3适合对窗口逻辑有特殊定制需求的场景。

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