Spark使用窗口函数collect_set获取薪资分桶全部下界全局数组的问题

问题解决思路 #

问题根源出在collect_set使用了带orderBy("salary")的窗口，Spark中排序窗口的默认帧范围是从数据集开头到当前行，所以只会逐步聚合当前行及之前的分桶下界，无法拿到全局所有分桶的下界值。

修改方案 #

只需要替换collect_set对应的窗口为全局无边界窗口即可，修改后的完整代码如下：

val spark = sparkSession
import spark.implicits._
    
val simpleData = Seq(("James", "Sales", 3000),
  ("Michael", "Sales", 3100),
  ("Robert", "Sales", 3200),
  ("Maria", "Finance", 3300),
  ("James", "Sales", 3400),
  ("Scott", "Finance", 3500),
  ("Jen", "Finance", 3600),
  ("Jeff", "Marketing", 3700),
  ("Kumar", "Marketing", 3800),
  ("Saif", "Sales", 3900)
)
val df = simpleData.toDF("employee_name", "department", "salary")
val windowSpec = Window.orderBy("salary")
val ntileFrame = df.withColumn("ntile", ntile(3).over(windowSpec))
val lowWindowSpec = Window.partitionBy("ntile")
val ntileMinDf = ntileFrame.withColumn("lower_bound", min("salary").over(lowWindowSpec))
// 新增全局无边界窗口
val globalWindow = Window.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)
val rangeDf = ntileMinDf.withColumn("range", collect_set("lower_bound").over(globalWindow))
rangeDf.show()

逻辑说明 #

Window.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)定义的全局窗口，会扫描整个数据集所有行的lower_bound值，去重后得到包含三个分桶下界的完整数组，所有行的range列都会拿到相同的全局数组，完全匹配预期输出。
如果需要数组按薪资大小有序排列，可以把collect_set("lower_bound")替换为sort_array(collect_set("lower_bound"))，得到的有序结果为[3000, 3400, 3700]。

内容的提问来源于stack exchange，提问作者pramodh