Pyspark 计数器字段，groupby 并递增 1

Question

我的数据框如下：

cola, colb
1, 2
1, 3
2, 1
2, 5

我要加一栏count：

cola, colb, count
1,     2,    1
1,     3,    2
2,     1,    1
2,     5,    2

需要对 colA 进行分组并将第一条记录的计数设置为 1，然后将每个连续的行递增 1。

我尝试使用窗口函数，但它对 colA 中的所有记录应用相同的计数而没有增量。

Answer 1

您已经发现窗口函数是可行的方法。可能你没用过rank函数。

import pyspark.sql.functions as F
from pyspark.sql import Window

l = [(1       , 2),
(1      , 3  ),
(1      , 2  ),
(2      , 1  ),
(2      , 5)]

columns = ['cola', 'colb']

df=spark.createDataFrame(l, columns)

w = Window.partitionBy('cola').orderBy('colb')

df = df.withColumn('count', F.rank().over(w))
df.show()

输出：

+----+----+-----+ 
|cola|colb|count| 
+----+----+-----+ 
|   1|   2|    1| 
|   1|   2|    1| 
|   1|   3|    3| 
|   2|   1|    1| 
|   2|   5|    2|
+----+----+-----+

如果您不想在相等的行之后出现间隙，您应该使用dense_rank 函数。