基於 DataFrame 中另一列的列的滾動總和

Question

我有一個如下所示的 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

從金額列中，我試圖根據日期列獲得 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 

Answer 1

根據建議，您可以在.rolling Date使用“ 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

Answer 2

我相信大熊貓的滾動方法是基於該指數的。 因此執行：

df.index = df['Date']

然后執行由您的時間范圍指定的滾動方法就可以解決問題。

另請參閱文檔（特別是文檔底部的文檔）： https : //pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.rolling.html

編輯：您也可以使用注釋中指出的on='Date'參數，無需重新編制索引。

Answer 3

這可以使用pandas_udf完成，並且看起來您想與“ ID”分組，因此我將其用作組ID。

spark = SparkSession.builder.appName('test').getOrCreate()
df = spark.createDataFrame([Row(ID=10001, d='2019-07-01', Amount=50),
                            Row(ID=10001, d='2019-05-01', Amount=15),
                            Row(ID=10001, d='2019-06-25', Amount=10),
                            Row(ID=10001, d='2019-05-27', Amount=20),
                            Row(ID=10002, d='2019-06-29', Amount=25),
                            Row(ID=10002, d='2019-07-18', Amount=35),
                            Row(ID=10002, d='2019-07-15', Amount=40)
                           ])
df = df.withColumn('date', F.to_date('d', 'yyyy-MM-dd'))
df = df.withColumn('prev_date', F.date_sub(df['date'], 28))
df.select(["ID", "prev_date", "date", "Amount"]).orderBy('date').show()
df = df.withColumn('amount_4wk_rolling', F.lit(0.0))
@pandas_udf(df.schema, PandasUDFType.GROUPED_MAP)
def roll_udf(pdf):
    for index, row in pdf.iterrows():
        d, pd = row['date'], row['prev_date']
        pdf.loc[pdf['date']==d, 'amount_4wk_rolling'] = np.sum(pdf.loc[(pdf['date']<=d)&(pdf['date']>=pd)]['Amount'])
    return pdf

df = df.groupby('ID').apply(roll_udf)
df.select(['ID', 'date', 'prev_date', 'Amount', 'amount_4wk_rolling']).orderBy(['ID', 'date']).show()

輸出：

+-----+----------+----------+------+
|   ID| prev_date|      date|Amount|
+-----+----------+----------+------+
|10001|2019-04-03|2019-05-01|    15|
|10001|2019-04-29|2019-05-27|    20|
|10001|2019-05-28|2019-06-25|    10|
|10002|2019-06-01|2019-06-29|    25|
|10001|2019-06-03|2019-07-01|    50|
|10002|2019-06-17|2019-07-15|    40|
|10002|2019-06-20|2019-07-18|    35|
+-----+----------+----------+------+

+-----+----------+----------+------+------------------+
|   ID|      date| prev_date|Amount|amount_4wk_rolling|
+-----+----------+----------+------+------------------+
|10001|2019-05-01|2019-04-03|    15|              15.0|
|10001|2019-05-27|2019-04-29|    20|              35.0|
|10001|2019-06-25|2019-05-28|    10|              10.0|
|10001|2019-07-01|2019-06-03|    50|              60.0|
|10002|2019-06-29|2019-06-01|    25|              25.0|
|10002|2019-07-15|2019-06-17|    40|              65.0|
|10002|2019-07-18|2019-06-20|    35|             100.0|
+-----+----------+----------+------+------------------+

Answer 4

對於pyspark，您可以只使用Window函數：sum + RangeBetween

from pyspark.sql import functions as F, Window

# skip code to initialize Spark session and dataframe

>>> df.show()
+-----+----------+------+
|   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.printSchema()
root
 |-- ID: long (nullable = true)
 |-- Date: string (nullable = true)
 |-- Amount: long (nullable = true)

win = Window.partitionBy('ID').orderBy(F.to_timestamp('Date').astype('long')).rangeBetween(-28*86400,0)

df_new = df.withColumn('amount_4wk_rolling', F.sum('Amount').over(win))

>>> df_new.show()
+------+-----+----------+------------------+
|Amount|   ID|      Date|amount_4wk_rolling|
+------+-----+----------+------------------+
|    25|10002|2019-06-29|                25|
|    40|10002|2019-07-15|                65|
|    35|10002|2019-07-18|               100|
|    15|10001|2019-05-01|                15|
|    20|10001|2019-05-27|                35|
|    10|10001|2019-06-25|                10|
|    50|10001|2019-07-01|                60|
+------+-----+----------+------------------+