在Spark Dataframe的窗口上创建组ID

Question

我有一个数据框，我想在每个Window分区中提供ID。例如我有

id | col |
1  |  a  |
2  |  a  |
3  |  b  |
4  |  c  |
5  |  c  |

所以我想要（基于与列col分组）

id | group |
1  |  1    |
2  |  1    |
3  |  2    |
4  |  3    |
5  |  3    |

我想使用窗口函数，但是无论如何我都找不到为每个窗口分配ID的方法。我需要类似的东西：

w = Window().partitionBy('col')
df = df.withColumn("group", id().over(w)) 

有什么办法可以达到这样的目标吗？（我不能简单地将col用作组ID，因为我有兴趣在多个列上创建一个窗口）

Answer 1

只需dense_rank 在Window函数上使用内置函数即可为您提供所需的结果

from pyspark.sql import window as W
import pyspark.sql.functions as f
df.select('id', f.dense_rank().over(W.Window.orderBy('col')).alias('group')).show(truncate=False)

这应该给你

+---+-----+
|id |group|
+---+-----+
|1  |1    |
|2  |1    |
|3  |2    |
|4  |3    |
|5  |3    |
+---+-----+