Spark/PySpark/Scala基于position列邻域范围计算每行value的降序排名

核心思路 #

该需求属于滑动邻域内的排名计算场景，可直接通过Spark的范围滑动窗口函数实现，无需自定义UDF，性能最优。
注意不要使用行偏移窗口rowsBetween，该窗口仅在position连续无缺失时生效，范围窗口rangeBetween无论position是否连续都能正确匹配「position差值绝对值不超过阈值」的邻域规则。

PySpark实现 #

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import rank, col

# 构造测试数据，已有dfTest可跳过本段
spark = SparkSession.builder.appName("neighbor_rank").getOrCreate()
data = [("a",1,8),("b",2,9),("c",3,10),("d",4,7),("e",5,9)]
dfTest = spark.createDataFrame(data, schema=["item","position","value"])

# 定义阈值和窗口规则
threshold = 1
window_spec = Window.orderBy(col("position")).rangeBetween(-threshold, threshold)

# 计算邻域内value降序排名
dfResult = dfTest.withColumn("rank", rank().over(window_spec.orderBy(col("value").desc())))

# 输出结果
dfResult.show()

Scala实现 #

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.rank

// 构造测试数据，已有dfTest可跳过本段
val data = Seq(("a",1,8),("b",2,9),("c",3,10),("d",4,7),("e",5,9))
val dfTest = spark.createDataFrame(data).toDF("item","position","value")

// 定义阈值和窗口规则
val threshold = 1
val windowSpec = Window.orderBy("position").rangeBetween(-threshold, threshold)

// 计算邻域内value降序排名
val dfResult = dfTest.withColumn("rank", rank().over(windowSpec.orderBy($"value".desc)))

// 输出结果
dfResult.show()

Spark SQL实现 #

-- 先将dfTest注册为临时视图
dfTest.createOrReplaceTempView("test_table")

-- 执行SQL查询得到结果
val dfResult = spark.sql("""
SELECT 
    item,
    position,
    value,
    RANK() OVER(ORDER BY position RANGE BETWEEN 1 PRECEDING AND 1 FOLLOWING ORDER BY value DESC) as rank
FROM test_table
""")

提示：如果要求相同value的排名相同且后续排名不跳号，将rank()替换为dense_rank()即可。

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