如何根据某些条件为 spark 数据框中的记录分配排名？

Question

给定一个数据框：

+-------+-------+
|   A   |   B   |
+-------+-------+
|      a|      1|
+-------+-------+
|      b|      2|
+-------+-------+
|      c|      5|
+-------+-------+
|      d|      7|
+-------+-------+
|      e|     11|
+-------+-------+    

我想根据条件为记录分配等级：

1. 排名从 1 开始
2. 如果（当前记录的 B - 先前记录的 B）为
3. 当（当前记录的 B - 先前记录的 B）> 2 时增加排名

所以我希望结果是这样的：

+-------+-------+------+
|   A   |   B   | rank |
+-------+-------+------+
|      a|      1|     1|
+-------+-------+------+
|      b|      2|     1|
+-------+-------+------+
|      c|      5|     2|
+-------+-------+------+
|      d|      7|     2|
+-------+-------+------+
|      e|     11|     3|
+-------+-------+------+

• spark 中的内置函数（如 rowNumber、rank、dense_rank）没有 提供任何功能来实现这一点。
• 我尝试使用全局变量 rank 并获取 使用滞后函数的先前记录值，但它没有给出 由于 Spark 中的分布式处理与 sql 不同，因此结果一致。

• 我尝试的另一种方法是将记录的滞后值传递给 UDF，同时生成新列并在 UDF 中应用条件。但我面临的问题是我可以获得列 A 和 B 的滞后值，但不能获得列排名。 这会产生错误，因为它无法解析列名排名：

  HiveContext.sql("SELECT df.*,LAG(df.rank, 1) OVER (ORDER BY B, 0) AS rank_lag, udfGetVisitNo(B,rank_lag) as rank FROM df")

• 我无法获取当前添加的列的滞后值。

• 另外我不想要需要使用 df.collect() 的方法，因为这个数据帧的大小非常大，并且在单个工作节点上收集它会导致内存错误。

还有其他方法可以达到同样的效果吗？ 我想知道一个时间复杂度为 O(n) 的解决方案，n 是记录数。

Answer 1

SQL 解决方案是

select a,b,1+sum(col) over(order by a) as rnk
from 
(
select t.*
,case when b - lag(b,1,b) over(order by a) <= 2 then 0 else 1 end as col
from t
) x

该解决方案假定排序基于列a。

SQL Server example