Пронумеровывать дубликаты в кадре данных Spark без перемешивания

Question

У меня есть фрейм данных Spark с данными временного ряда (упорядоченный по индексу) и столбцом значений. Это выглядит так:

idx    value    counter
0      2        1
1      2        2
2      2        3
3      1        1
4      1        2
5      0        1
6      0        2
7      0        3
8      0        4
9      1        1
10     1        2

Столбец "counter" - это то, что я ищу. Когда он видит новое значение, счетчик должен сброситься до 1 и начать отсчет, пока значение не изменится.

Я знаю, что можно назначить groupID для каждой группы последовательных дубликатов, а затем вызвать

df.withColumn('counter', row_number().over(Window.partitionBy('groupID').orderBy('idx')))

Однако это может привести к случайному перемешиванию, поскольку все элементы уже разделены по мере необходимости. Я просто хочу счетчик, который сбрасывается, когда он видит новое значение. Есть ли способ добиться этого с моим текущим разделением или сказать Spark, что он должен работать с группами последовательных дубликатов, которые я определил, но без фактического перемещения данных. Проблема в том, что Window.partitionBy (col) всегда перемещает данные.

Answer 1

Конечно, это можно сделать:

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

partition = Window.partitionBy("groupID").orderBy("idx")
df = df.withColumn("counter", F.count(F.col("value")).over(partition))

df.show(10, False)