如何在 Pandas 数据帧中使用 collect_list 复制 Spark window 函数？

Question

我有一个初始数据框

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 我使用 collect_list 函数使用 window 函数对数据框进行编码，同时考虑到分组列 'A' 并考虑到列 'B'排序以创建具有累积列表的列

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

df1.orderBy('A','B').show()

+---+---+---+------------------------+
|A  |B  |C  |D                       |
+---+---+---+------------------------+
|1  |1  |10 |[10]                    |
|1  |2  |11 |[10, 11, 12]            |
|1  |2  |12 |[10, 11, 12]            |
|2  |1  |14 |[14, 15, 16]            |
|2  |1  |15 |[14, 15, 16]            |
|2  |1  |16 |[14, 15, 16]            |
|3  |1  |13 |[13]                    |
|4  |1  |17 |[17]                    |
|4  |2  |18 |[17, 18]                |
|4  |3  |19 |[17, 18, 19]            |
|4  |4  |19 |[17, 18, 19, 19]        |
|4  |5  |20 |[17, 18, 19, 19, 20, 20]|
|4  |5  |20 |[17, 18, 19, 19, 20, 20]|
+---+---+---+------------------------+

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

我尝试使用一些“普通”python 代码，但可能有更直接的方法。

Answer 1

在 pandas 中解决此问题的一种方法是使用两个 groupby，即首先在列 A 上对数据框进行分组，然后为每个组应用自定义函数 collect_list，其中按列 B 转换组输入并使用 list

累积聚合列 C

def collect_list(g):
    return g.groupby('B')['C'].agg(list).cumsum()

df.sort_values(['A', 'B']).merge(
    df.groupby('A').apply(collect_list).reset_index(name='D'))

    A  B   C                         D
0   1  1  10                      [10]
1   1  2  11              [10, 11, 12]
2   1  2  12              [10, 11, 12]
4   2  1  14              [14, 15, 16]
5   2  1  15              [14, 15, 16]
6   2  1  16              [14, 15, 16]
3   3  1  13                      [13]
7   4  1  17                      [17]
8   4  2  18                  [17, 18]
9   4  3  19              [17, 18, 19]
10  4  4  19          [17, 18, 19, 19]
11  4  5  20  [17, 18, 19, 19, 20, 20]
12  4  5  20  [17, 18, 19, 19, 20, 20]

如何在 Pandas 数据帧中使用 collect_list 复制 Spark window 函数？

How to replicate a Spark window function with collect_list in Pandas dataframes?

pandas

apache-spark

pyspark

window