Spark-窗口函数实现原理及各种写法

link之家
链接快照平台
输入网页链接，自动生成快照
标签化管理网页链接
1 窗口函数能做什么？

在日常的数据操作中，除了普通的单条顺序处理，就是聚合操作了，其实分别代表了最基本的非Shuffle和Shuffle操作。
但实际上，还有一种，就是处理当前这条数据时，还需要考虑上下文，比如随机取N条，那我必须知道是不是取够N条了，当前这条是否要丢弃。再比如移动平均，需要用到前后数据参与计算平均值。
Spark中，最基本的窗口函数可以这么写：
select id,category,Row_number() over(partition by category order by id) rank from t where rank <=10;
这是一个典型的TOPN操作，取每个category分类中最小的10个id对应的行。
这条sql中包含了两个部分：
窗口定义部分，即 partition by category
Window函数部分，即 order by id
分别定义了如何切分窗口，以及在每个窗口中，如何做数据处理。
那么在Spark中，它的物理计划是什么样的呢？
== Physical Plan ==
Window [row_number() windowspecdefinition(category#8, id#7 ASC NULLS FIRST, specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$())) AS rank#11], [category#8], [id#7 ASC NULLS FIRST]
+- *(1) Sort [category#8 ASC NULLS FIRST, id#7 ASC NULLS FIRST], false, 0
   +- Exchange hashpartitioning(category#8, 100), ENSURE_REQUIREMENTS, [id=#10]
      +- LocalTableScan [id#7, category#8]
整体逻辑为：
按照category做repartition
每个partition内，按照(category,id)进行排序
执行窗口函数，即输出3个字段
2 窗口选取部分
窗口部分有3部分
partition by 必选
order by 可选(row_number和range between是必选的)
rows between 、range between
先创建一个测试表：
 > spark.createDataset(Seq((1, "b"), (2, "c"), (3, "a"), (4, "a"), (3, "c"),(4, "c"))).toDF("id", "category")
+---+--------+
|id |category|
+---+--------+
|1  |b       |
|2  |c       |
|3  |a       |
|4  |a       |
|3  |c       |
|4  |c       |
+---+--------+
2.1 row between 按照距离取数
2.1.1.rows between unbounded preceding and current row
窗口内，从开始一直取到当前数据
> select id,category,sum(id) over(partition by category rows between unbounded preceding and current row) rank from t
+---+--------+----+
|id |category|rank|
+---+--------+----+
|2  |c       |2   |
|3  |c       |5   |
|4  |c       |9   |
|1  |b       |1   |
|3  |a       |3   |
|4  |a       |7   |
+---+--------+----+
可以看到rank值是id的从开始当每一行当前id的累加值。
2.1.2.rows between 2 preceding and 2 following
取当前行的前两条及后两条来参与计算，一般用于移动平均值
> select id,category,sum(id) over(partition by category rows between 1 preceding and 1 following) rank from t
+---+--------+----+
|id |category|rank|
+---+--------+----+
|2  |c       |5   |
|3  |c       |9   |
|4  |c       |7   |
|1  |b       |1   |
|3  |a       |7   |
|4  |a       |7   |
+---+--------+----+
可以看到每个rank值都是取了前后一条的数据进行sum计算，如果当前行向前或向后取不到对应的条数，那么就放弃。
2.2 range between 按照大小取数
RANGE BETWEEN 这种就是以当前值为锚点进行计算。

比如RANGE BETWEEN 4 PRECEDING AND 7 FOLLOWING

当前值为10的话就去前后的值在6到17之间的数据。
> select id,category,sum(id) over(partition by category order by id range between 0 preceding and 1 following) rank from t
+---+--------+----+
|id |category|rank|
+---+--------+----+
|2  |c       |5   |
|3  |c       |7   |
|4  |c       |4   |
|1  |b       |1   |
|3  |a       |7   |
|4  |a       |4   |
+---+--------+----+
sql的含义是取当前值之前相等的，之后差为1的。比如当前值为3，那么取id在[2,4]之间的数据，注意向前或向后是可以取多条数据的。
另外，order by好像是必须的，并会指定该字段为之后取数的判断字段。
2.3 总结
几种窗口函数其实是可以相互组合的。比如：
> select id,category,sum(id) over(partition by category order by id range between 1 preceding and current row) rank from t
+---+--------+----+
|id |category|rank|
+---+--------+----+
|2  |c       |5   |
|3  |c       |9   |
|4  |c       |7   |
|1  |b       |1   |
|3  |a       |7   |
|4  |a       |7   |
+---+--------+----+
可以非常方便的满足窗口计算的要求。
3 窗口数据部分
该部分就是在窗口定义好，即窗口内的数据已经获取完毕后，对这些数据的处理方式。
具体有3种：
聚合函数：count\sum\avg等
排序函数：row_number、rank
分析函数：cume_dist函数计算当前值在窗口中的百分位数
row_number函数即为当前行标记行数，从1开始，必须包含order by部分。
4 DataFrame写法
前半部分都是用Sql来举的例子，那么DataFrame是怎么调用窗口函数呢。
import org.apache.spark.sql.expressions.{Window, WindowSpec}
val window: WindowSpec = Window.partitionBy("category")
  .orderBy("id")
  .rangeBetween(0l, 1l)
df.withColumn("rank",sum("id").over(window))
  .show(false)
+---+--------+----+
|id |category|rank|
+---+--------+----+
|2  |c       |5   |
|3  |c       |7   |
|4  |c       |4   |
|1  |b       |1   |
|3  |a       |7   |
|4  |a       |4   |
+---+--------+----+
这段代码与下列sql的结果完全一致，即上文2.2 部分。
 select id,category,sum(id) over(partition by category order by id range between 0 preceding and 1 following) rank from t
正如上述例子中，窗口函数主要包含两个部分：

指定窗口特征（wSpec）：
“partitionyBY” 定义数据如何分组；在上面的例子中，是用户 site
“orderBy” 定义分组中的排序
“rowsBetween” 定义窗口的大小
指定窗口函数函数：

指定窗口函数函数，你可以使用 org.apache.spark.sql.functions 的“聚合函数（Aggregate Functions）”和”窗口函数（Window Functions）“类别下的函数
4.1 累计汇总
val wSpac: WindowSpec = Window.partitionBy("site")
    .orderBy("date")
    .rowsBetween(Long.MinValue, 1)
df.withColumn("cumSum",sum(df("user_cnt")).over(wSpec)).show()
rowsBetween(Long.MinValue, 0) #窗口的大小是按照排序从最小值到当前行

rangeBetween(currentRow(), lit(1)) #表示当前行到下一行
4.2 前一行数据
val wSpac: WindowSpec = Window.partitionBy("site")
    .orderBy("date")
df.withColumn("cumSum",
lag(df("user_cnt"),1).over(wSpec)).show()
lag(field, n): 就是取从当前字段往前第n个值，这里是取前一行的值
4.3 排名
val wSpec=Window.partitionBy("site")
    .orderBy("date")
df.withColumn("rank",rank().over(wSpec)).show()
rank()在提取每个分组的前N项特别有用。