基于DataFrame中另一列的列的滚动总和

问题描述

我有一个如下所示的DataFrame

 ID      Date      Amount   

10001   2019-07-01   50     
10001   2019-05-01   15
10001   2019-06-25   10   
10001   2019-05-27   20
10002   2019-06-29   25
10002   2019-07-18   35
10002   2019-07-15   40

从Amount列中,我尝试根据Date列获得4周累计总和。我的意思是,基本上我还需要一个列(比方说Amount_4wk_Rolling),它将有一个返回4周的所有行的金额列的总和。因此,如果行中的日期是2019-07-01,则Amount_4wk_Rolling列值应该是日期介于2019-07-01和2019-06-04(2019-07-01减去28天)之间的所有行的金额之和。 因此,新的DataFrame应该如下所示。

 ID        Date      Amount  amount_4wk_rolling
10001   2019-07-01    50       60
10001   2019-05-01    15       15
10001   2019-06-25    10       30
10001   2019-05-27    20       35
10002   2019-06-29    25       25
10002   2019-07-18    35       100
10002   2019-07-15    40       65

我尝试使用窗口函数,但它不允许我根据特定列的值选择窗口

Edit:
 My data is huge...about a TB in size. Ideally, I would like to do this in spark rather that in pandas 

解决方案

按照建议,您可以将.rollingOn[2-0]>与"28d"一起使用。

从您的示例值来看,您似乎还希望滚动窗口按ID分组。

试试:

import pandas as pd
from io import StringIO

s = """
 ID      Date      Amount   

10001   2019-07-01   50     
10001   2019-05-01   15
10001   2019-06-25   10   
10001   2019-05-27   20
10002   2019-06-29   25
10002   2019-07-18   35
10002   2019-07-15   40
"""

df = pd.read_csv(StringIO(s), sep="s+")
df['Date'] = pd.to_datetime(df['Date'])
amounts = df.groupby(["ID"]).apply(lambda g: g.sort_values('Date').rolling('28d', on='Date').sum())
df['amount_4wk_rolling'] = df["Date"].map(amounts.set_index('Date')['Amount'])
print(df)

输出:

      ID       Date  Amount  amount_4wk_rolling
0  10001 2019-07-01      50                60.0
1  10001 2019-05-01      15                15.0
2  10001 2019-06-25      10                10.0
3  10001 2019-05-27      20                35.0
4  10002 2019-06-29      25                25.0
5  10002 2019-07-18      35               100.0
6  10002 2019-07-15      40                65.0

相关文章