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

2019年7月28日 186次阅读

给定一个数据帧：

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

我想根据条件为记录分配排名：

>以1开始排名
>如果(当前记录的B – 先前记录的B)<= 2,则指定等级=先前记录的等级
>当(当前记录的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这样的内置函数没有
提供实现此目的的任何功能.
>我尝试使用全局变量排名和提取
以前的记录值使用滞后函数,但它没有给出
与sql不同,由于Spark中的分布式处理,结果一致.
>我尝试的另一种方法是在生成新列并在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是记录的编号.

最佳答案 一个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

点赞