汇总为pyspark中3个最大值的总和

Question

我有一个按“ id”和“ type”分组的数据框：

+---+----+-----+
| id|type|count|
+---+----+-----+
|  0|   A|    2|
|  0|   B|    3|
|  0|   C|    1|
|  0|   D|    3|
|  0|   G|    1|
|  1|   A|    0|
|  1|   C|    1|
|  1|   D|    1|
|  1|   G|    2|
+---+----+-----+

我现在想按“ id”分组并获得3个最大值的总和：

+---+-----+
| id|count|
+---+-----+
|  0|    8|
|  1|    4|
+---+-----+

我如何在pyspark中做到这一点，以便计算相对高效？

找到了解决方案here

Answer 1

您可以使用以下代码执行此操作

from pyspark.sql.functions import *
from pyspark.sql.window import Window
df = spark.createDataFrame([
(0, "A", 2),
(0,"B", 3),
(0,"C", 1),
(0,"D", 3),
(1,"A", 0),
(1,"C", 1),
(1,"D", 1),
(1,"G", 2)
], ("id", "type", "count"))

my_window = Window.partitionBy("id").orderBy("count")
df.withColumn("last_3", lead("count").over(my_window)).groupBy("id").agg(sum("last_3")).show()

输出：

+---+-----------+
| id|sum(last_3)|
+---+-----------+
|  0|          8|
|  1|          4|
+---+-----------+

详细信息：窗口按id划分数据并按count对其进行排序，然后创建一个新列，其中lead使用此窗口并返回该组中的下一个值（ (0,C,1)是id=0组中最低的元组，因此它接收到的值是2，因为它是该组中第二高的元组（来自元组(0,A,2)等等，最高的元组没有以下值，而是被分配了null。最后执行分组运算和求和。