Поле счетчика Pyspark, групповое и увеличение на 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|
+----+----+-----+

Если вам не нужны пробелы после одинаковых строк, вам следуетиспользуйте функцию dens_rank .