Pandas滚动窗口与asof merge对比：按标识回溯N天取最新uniqueid

问题描述 #

需要按identifier分组后在滚动窗口执行计算，不确定用双向asof merge还是滚动窗口函数，具体问题如下：

数据结构 #

>>> df
           identifier  uniqueid
2021-08-01          A       0.0
2021-08-02          B       1.0
2021-08-03          C       2.0
2021-08-04          D       3.0
2021-08-05        NaN       NaN
...               ...       ...
2022-05-23          Z     295.0
2022-05-24         AA     296.0
2022-05-25         AB     297.0
2022-05-26         AC     298.0
2022-05-27        NaN       NaN

[300 rows x 2 columns]

可复现数据代码 #

import datetime
import string
import pandas as pd
import numpy as np

np.random.seed(100)

idx = pd.date_range(start=datetime.datetime(2021,8,1), periods=300, freq="D")
names =  ([x for x in string.ascii_uppercase] + ["AA", "AB", "AC", "AD"])*10
uniqueid = range(300)

df = pd.DataFrame(
{
  "identifier" : names,
  "uniqueid":  uniqueid,
}, 
index=idx
)

censor = np.random.random(300) < 0.2
df.iloc[censor, :] = np.nan

需求说明 #

对每一天的每一个identifier，回溯过去N天，获取该identifier对应的最近一条uniqueid值，若最近记录超出N天窗口范围则不计入结果。以N=30为例，期望输出格式如下：

2021-08-01           []
2021-08-02           []
...
2021-08-31           [(A, 30), (B, 1),  (C, 2), (D, 3), (G,6).... # 注意E、F为NA
2021-09-01           [(A, 30), (B, 31), (C, 2), ....
2021-09-02           [(A, 30), (B, 31), (D, 3), ....  # 注意9月2日C为NA，且8月3日记录已经超出30天窗口

遇到的问题 #

尝试用如下代码实现：

df.groupby("identifier").rolling(30).agg(lambda x: x.iloc[-1])

运行时报错：

ValueError: Length of passed values is 244, index implies 300.

需要明确应该选用asof merge还是滚动窗口实现需求，以及如何解决上述报错。

解决方案 #

报错的原因是groupby时会自动跳过identifier为NaN的行，分组后每组的长度和原索引长度不匹配，直接调用rolling聚合后返回的结果长度和原DataFrame不一致才触发了值错误。

实现该需求用滚动窗口+分组的方式更直接，修改思路如下：

1. 先删除identifier为NaN的行，避免分组时的空值干扰
2. 对分组后的时间索引设置时间偏移窗口，而非固定30行的窗口，避免中间有空值时窗口计算错误
3. 每组取滚动窗口的最后一个有效值，再按日期重新聚合得到每日的元组列表

完整可运行代码如下：

N = 30
# 先清理空值行
df_clean = df.dropna(subset=['identifier']).copy()
# 分组后用时间窗口滚动，取最近的uniqueid
df_clean['last_uniqueid'] = df_clean.groupby('identifier')\
    .rolling(window=f'{N}D', closed='both')['uniqueid']\
    .apply(lambda x: x.iloc[-1])\
    .reset_index(level=0, drop=True)

# 按日期聚合，得到每日的(identifier, last_uniqueid)列表
result = df_clean.groupby(df_clean.index)\
    .apply(lambda x: list(zip(x['identifier'], x['last_uniqueid'])))\
    .reindex(df.index, fill_value=[])

# 输出前40天结果验证
print(result.head(40))

如果用asof merge实现，需要先生成所有日期和所有identifier的笛卡尔积，再每个日期关联最近30天内的有效记录，实现复杂度比滚动窗口高，性能也不如上述方案。

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