如何使用累积和在 Pandas Dataframe 中执行窗口函数操作？

Posted 2023-03-11

【中文标题】如何使用累积和在 Pandas Dataframe 中执行窗口函数操作？

【英文标题】：How to perform a Window Function operation in a Pandas Dataframe using a cumulative sum?

【发布时间】：2021-12-13 00:13:11

【问题描述】：

我有一个初始数据框

df1 = 
+---+---+---+
|  A|  B|  C|
+---+---+---+
|  1|  1| 10|
|  1|  2| 11|
|  1|  2| 12|
|  3|  1| 13|
|  2|  1| 14|
|  2|  1| 15|
|  2|  1| 16|
|  4|  1| 17|
|  4|  2| 18|
|  4|  3| 19|
|  4|  4| 19|
|  4|  5| 20|
|  4|  5| 20|
+---+---+---+

使用 pyspark，我使用一个窗口函数对数据帧进行编码，其中考虑了“A”列并考虑了“B”列的排序。

spec = Window.partitionBy('A').orderBy('B')
df1 = df1.withColumn('D',sum('C').over(spec))

df1.show()

+---+---+---+-----+
|  A|  B|  C|    D|
+---+---+---+-----+
|  1|  1| 10| 10.0|
|  1|  2| 11| 33.0|
|  1|  2| 12| 33.0|
|  2|  1| 14| 45.0|
|  2|  1| 15| 45.0|
|  2|  1| 16| 45.0|
|  3|  1| 13| 13.0|
|  4|  1| 17| 17.0|
|  4|  2| 18| 35.0|
|  4|  3| 19| 54.0|
|  4|  4| 19| 73.0|
|  4|  5| 20|113.0|
|  4|  5| 20|113.0|
+---+---+---+-----+

是否可以使用 Pandas Dataframe 进行相同的计算？

我尝试过使用

df['D'] = df.sort_values(['A','B']).groupby(['A', 'B'])['C'].transform('cumsum')

但结果不一样

谢谢

【参考方案1】：

你可以sort，然后cumsum在'A'中，然后groupby + max在['A'，'B']组中使用transform，这样你就可以分配回结果。

df1['D'] = (df1.sort_values(['A', 'B'])
               .groupby('A')['C'].cumsum()
               .groupby([df1['A'], df1['B']]).transform('max'))

---

    A  B   C    D
0   1  1  10   10
1   1  2  11   33
2   1  2  12   33
3   3  1  13   13
4   2  1  14   45
5   2  1  15   45
6   2  1  16   45
7   4  1  17   17
8   4  2  18   35
9   4  3  19   54
10  4  4  19   73
11  4  5  20  113
12  4  5  20  113

【参考方案2】：

在 pandas 中，我们可以在A 和B 上使用groupby sum。然后groupby cumsum 刚刚超过A。要将结果添加回 DataFrame rename 到新列名，然后 join 将结果返回到初始组键 ['A', 'B'] 上的 DataFrame：

df1 = df1.join(
    df1.groupby(by=['A', 'B'])['C'].sum()
        .groupby(level='A').cumsum()
        .rename('D'),
    on=['A', 'B']
)

df1:

    A  B   C    D
0   1  1  10   10
1   1  2  11   33
2   1  2  12   33
3   3  1  13   13
4   2  1  14   45
5   2  1  15   45
6   2  1  16   45
7   4  1  17   17
8   4  2  18   35
9   4  3  19   54
10  4  4  19   73
11  4  5  20  113
12  4  5  20  113

---

设置：

import pandas as pd

df1 = pd.DataFrame(
    'A': [1, 1, 1, 3, 2, 2, 2, 4, 4, 4, 4, 4, 4],
    'B': [1, 2, 2, 1, 1, 1, 1, 1, 2, 3, 4, 5, 5],
    'C': [10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 19, 20, 20]
)