如何根據某些條件為火花數據框中的記錄分配排名？

Question

給定一個數據幀：

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

我想根據條件為記錄分配排名：

1. 以1開始排名
2. 如果（當前記錄的B - 先前記錄的B）<= 2，則分配等級=先前記錄的等級
3. 當前記錄的B（前一記錄的B）> 2時的增量等級

所以我希望結果是這樣的：

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

• 像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