Накопительная сумма в писпарке

Question

Я пытаюсь вычислить совокупную сумму на класс. Код работает нормально, используя sum (df.value) .over (Window.partitionBy ('class'). OrderBy ('time'))

df = sqlContext.createDataFrame( [(1,10,"a"),(3,2,"a"),(1,2,"b"),(2,5,"a"),(2,1,"b"),(9,0,"b"),(4,1,"b"),(7,8,"a"),(3,8,"b"),(2,5,"a"),(0,0,"a"),(4,3,"a")], 
                                     ["time", "value", "class"] )

time|value|class|
+----+-----+-----+
|   1|   10|    a|
|   3|    2|    a|
|   1|    2|    b|
|   2|    5|    a|
|   2|    1|    b|
|   9|    0|    b|
|   4|    1|    b|
|   7|    8|    a|
|   3|    8|    b|
|   2|    5|    a|
|   0|    0|    a|
|   4|    3|    a|


df.withColumn('cumsum_value', sum(df.value).over(Window.partitionBy('class').orderBy('time'))).show()


time|value|class|cumsum_value|
+----+-----+-----+------------+
|   1|    2|    b|           2|
|   2|    1|    b|           3|
|   3|    8|    b|          11|
|   4|    1|    b|          12|
|   9|    0|    b|          12|
|   0|    0|    a|           0|
|   1|   10|    a|          10|
|   2|    5|    a|          20|
|   2|    5|    a|          20|
|   3|    2|    a|          22|
|   4|    3|    a|          25|
|   7|    8|    a|          33|
+----+-----+-----+------------+

, но он не работает с дублирующимися строками. Желаемый результат должен быть:

 time|value|class|cumsum_value|
+----+-----+-----+------------+
|   1|    2|    b|           2|
|   2|    1|    b|           3|
|   3|    8|    b|          11|
|   4|    1|    b|          12|
|   9|    0|    b|          12|
|   0|    0|    a|           0|
|   1|   10|    a|          10|
|   2|    5|    a|          15|
|   2|    5|    a|          20|
|   3|    2|    a|          22|
|   4|    3|    a|          25|
|   7|    8|    a|          33|
+----+-----+-----+------------+

Answer 1

Добавляя к комментарию @pault, я бы предложил вычисление row_number() на основе orderBy('time', 'value') и затем использовал этот столбец в orderBy другого окна (w2), чтобы получить cum_sum.

Это будет обрабатывать оба случая, когда время одинаково и значение одинаково, а время одинаково, но значение не равно.

from pyspark.sql import functions as F
from pyspark.sql.window import Window
w1=Window().partitionBy("class").orderBy("time","value")
w2=Window().partitionBy("class").orderBy('rownum')
df.withColumn('rownum', F.row_number().over(w1))\
  .withColumn('cumsum_value', F.sum("value").over(w2)).drop('rownum').show()

+----+-----+-----+------------+
|time|value|class|cumsum_value|
+----+-----+-----+------------+
|   1|    2|    b|           2|
|   2|    1|    b|           3|
|   3|    8|    b|          11|
|   4|    1|    b|          12|
|   9|    0|    b|          12|
|   0|    0|    a|           0|
|   1|   10|    a|          10|
|   2|    5|    a|          15|
|   2|    5|    a|          20|
|   3|    2|    a|          22|
|   4|    3|    a|          25|
|   7|    8|    a|          33|
+----+-----+-----+------------+