pyspark--Window 特定统计值所在行保留

假定有如下数据：

from pyspark.sql.window import Window
from pyspark.sql import functions as F

col_names = ["name", "date", "score"]
value = [
    ("Ali", "2020-01-01", 10.0),
    ("Ali", "2020-01-02", 15.0),
    ("Ali", "2020-01-03", 20.0),
    ("Ali", "2020-01-04", 25.0),
    ("Ali", "2020-01-05", 30.0),
    ("Bob", "2020-01-01", 15.0),
    ("Bob", "2020-01-02", 20.0),
    ("Bob", "2020-01-03", 30.0)
]
df = spark.createDataFrame(value, col_names)

现在要统计获取每个人最大的分数，并且需要保留该数值所在行的所有数据，那么用常规的groupby我们会得到如下结果：

df.groupby("name").max("score").show()
+----+----------+
|name|max(score)|
+----+----------+
| Bob|      30.0|
| Ali|      30.0|
+----+----------+

很明显这个结果不是我们想要的，而如果在此基础上要保留其他信息就需要用join获取，这样无疑大大增加了计算量。那么有没有一种既能满足需求，又能减少计算量的方法呢？

答案必然是有的，这时候就该Window出场了。

我们建一个以name分组的窗对象，并在分组内对score进行倒排序，然后按照排序给组内的数据添加一个排名，这样每个人的分数就按照从大到小进行了排序，第一名就是最大分数。具体代码如下：

w = Window.partitionBy("name").orderBy(F.desc("score"))
df_rank = df.withColumn("rank",F.rank().over(w))
df_rank.show()

+----+----------+-----+----+
|name|      date|score|rank|
+----+----------+-----+----+
| Bob|2020-01-03| 30.0|   1|
| Bob|2020-01-02| 20.0|   2|
| Bob|2020-01-01| 15.0|   3|
| Ali|2020-01-05| 30.0|   1|
| Ali|2020-01-04| 25.0|   2|
| Ali|2020-01-03| 20.0|   3|
| Ali|2020-01-02| 15.0|   4|
| Ali|2020-01-01| 10.0|   5|
+----+----------+-----+----+

分析结果可以看出来每个人的rank都是按照其分数从大到小进行排序的，这样要保留最大分数只需要过滤rank=1的行即可。需要注意的是rank()和dense_rank()函数的区别，dense_rank()中有并列排名时，下一名次不受影响。完整代码如下：

from pyspark.sql.window import Window
from pyspark.sql import functions as F

col_names = ["name", "date", "score"]
value = [
    ("Ali", "2020-01-01", 10.0),
    ("Ali", "2020-01-02", 15.0),
    ("Ali", "2020-01-03", 20.0),
    ("Ali", "2020-01-04", 25.0),
    ("Ali", "2020-01-05", 30.0),
    ("Bob", "2020-01-01", 15.0),
    ("Bob", "2020-01-02", 20.0),
    ("Bob", "2020-01-03", 30.0)
]
df = spark.createDataFrame(value, col_names)
df.show()

df.groupby("name").max("score").show()

w = Window.partitionBy("name").orderBy(F.desc("score"))
df_rank = df.withColumn("rank",F.rank().over(w))
df_rank.show()

df_rank.filter(F.col("rank")==1).show()

还有一种更加直接的方式，直接对分组求最大值，再过滤分数等于最大值的行即可。其本质与前一种方法一样，代码及结果如下：

w = Window.partitionBy("name").orderBy(F.desc("score"))
df_rank = df.withColumn("max",F.max("score").over(w))
df_rank.show()

df_rank.filter(F.col("score")==F.col("max")).show()

+----+----------+-----+----+
|name|      date|score| max|
+----+----------+-----+----+
| Bob|2020-01-03| 30.0|30.0|
| Bob|2020-01-02| 20.0|30.0|
| Bob|2020-01-01| 15.0|30.0|
| Ali|2020-01-05| 30.0|30.0|
| Ali|2020-01-04| 25.0|30.0|
| Ali|2020-01-03| 20.0|30.0|
| Ali|2020-01-02| 15.0|30.0|
| Ali|2020-01-01| 10.0|30.0|
+----+----------+-----+----+

+----+----------+-----+----+
|name|      date|score| max|
+----+----------+-----+----+
| Bob|2020-01-03| 30.0|30.0|
| Ali|2020-01-05| 30.0|30.0|
+----+----------+-----+----+

参考链接：pyspark-groupby-and-max-value-selection