如何在datetime索引的基础上高效滚动分组

Question

让我们首先创建一个中等大小的随机df:

size_now = 1e5
df = pd.DataFrame({"time":np.random.randint(1000,size=int(size_now)),
"id":np.random.randint(1000,size=int(size_now)),
"data":np.random.randint(1000,size=int(size_now))})
df.time =  pd.to_datetime(df.time,unit="s")

df看起来像这样：

print(df)
executed in 9ms, finished 15:16:05 2021-07-19
time   id  data
0     1970-01-01 00:15:01  915   299
1     1970-01-01 00:08:30   31   940
2     1970-01-01 00:02:55  925   600
3     1970-01-01 00:12:48  554   935
4     1970-01-01 00:07:43  565   638
...                   ...  ...   ...
99995 1970-01-01 00:16:27  707   615
99996 1970-01-01 00:01:06  298    38
99997 1970-01-01 00:14:26   37    43
99998 1970-01-01 00:16:09   66   791
99999 1970-01-01 00:01:45  376   854

我意识到滚动和可以很快：

df.set_index("time").sort_index().rolling("2500ms")["data"].sum()
time
1970-01-01 00:00:00        66.0
1970-01-01 00:00:00       809.0
1970-01-01 00:00:00       879.0
1970-01-01 00:00:00      1329.0
1970-01-01 00:00:00      1729.0
...   
1970-01-01 00:16:39    124998.0
1970-01-01 00:16:39    125265.0
1970-01-01 00:16:39    126050.0
1970-01-01 00:16:39    126782.0
1970-01-01 00:16:39    127766.0
Name: data, Length: 100000, dtype: float64

只花了0.031秒就完成了任务。然而，我需要的是按id分组并按时执行滚动。所以我的目标是：

df.groupby("id").rolling("2500ms",on="time")["data"].sum()
id   time               
0    1970-01-01 00:05:43      161.0
1970-01-01 00:12:05      288.0
1970-01-01 00:00:35      981.0
1970-01-01 00:00:22      479.0
1970-01-01 00:09:33      834.0
...   
999  1970-01-01 00:15:07    28159.0
1970-01-01 00:06:39    29035.0
1970-01-01 00:06:15    29538.0
1970-01-01 00:09:07    29788.0
1970-01-01 00:00:48    30736.0
Name: data, Length: 100000, dtype: float64

但完成任务花了2秒。我想知道是否有更有效的方法。

我认为预分类可能会有所帮助，所以我尝试了：

df = df.set_index(["id","time"]).sort_index()

这也很快。在这之后，我们只需要将数据帧切割成多个小帧，然后进行滚动求和并将它们连接在一起。。。所以我尝试了

df.groupby(level=0).rolling("2500ms")

但它给了我这个错误

ValueError：窗口必须是整数

所以我尝试了：

df.groupby(level=0).rolling("2500ms",on=df.index.levels[1]).sum()

这打破了我的木星笔记本，给了我内核死错误。

所以我尝试了

df_tmp_list = []
for j in df.index.unique(level=0):
df_tmp = df.loc[j]
df_tmp = df_tmp.rolling("2500ms").sum()
df_tmp["id"] = j
df_tmp_list.append(df_tmp)
#     break
pd.concat(df_tmp_list).reset_index().set_index(["id","time"])

这也花了2.6秒…对我来说，最有效的方法是什么？

df.groupby("id").rolling("2500ms",on="time")["data"].sum()

编辑感谢大家的回答。似乎每个人都比我跑得快。我需要2秒，其他人需要50秒。。。确切地说，完整的代码是：

import pandas as pd
import numpy as np
size_now = 1e5
df = pd.DataFrame({"time":np.random.randint(1000,size=int(size_now)),
"id":np.random.randint(1000,size=int(size_now)),
"data":np.random.randint(1000,size=int(size_now))})
df.time =  pd.to_datetime(df.time,unit="s")
%time df.groupby("id").rolling("2500ms",on="time")["data"].sum()

我的输出是

CPU times: user 2.3 s, sys: 24 ms, total: 2.32 s
Wall time: 2.32 s

我实际上是在AWS-EC2上运行的，有4个CPU和16GB RAM。我尝试过更新包：

510  conda update pandas
511  conda update numpy

所以速度没有提高。。。有什么帮助吗？我很沮丧。。。

我试着用我的MacBook Pro，2.8 GHz四核英特尔酷睿i7和16 GB 1600 MHz DDR3。它花了450毫秒。。。仍然比你们所有人慢得多。。。。请帮帮我…

Edit2，所以这是由于一些神秘的安装问题。我必须用最新的下载更新anaconda。

523  curl -O https://repo.anaconda.com/archive/Anaconda3-2021.05-Linux-x86_64.sh
525  bash Anaconda3-2021.05-Linux-x86_64.sh -u

现在它以50毫秒的速度运行。问题解决了！

Answer 1

我认为你可以获得任何东西：

>>> %timeit df.set_index('time').groupby('id').rolling('2500ms').sum()
57.9 ms ± 1.14 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
>>> %timeit df.groupby("id").rolling("2500ms", on="time")["data"].sum()
33.9 ms ± 466 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Answer 2

Edit2，所以这是由于一些神秘的安装问题。我必须用最新的下载更新anaconda。

523  curl -O https://repo.anaconda.com/archive/Anaconda3-2021.05-Linux-x86_64.sh
525  bash Anaconda3-2021.05-Linux-x86_64.sh -u

现在它以50毫秒的速度运行。问题解决了！