Eu tenho o seguinte quadro de dados mostrando a receita das compras.

+-------+--------+-------+
|user_id|visit_id|revenue|
+-------+--------+-------+
|      1|       1|      0|
|      1|       2|      0|
|      1|       3|      0|
|      1|       4|    100|
|      1|       5|      0|
|      1|       6|      0|
|      1|       7|    200|
|      1|       8|      0|
|      1|       9|     10|
+-------+--------+-------+

Finalmente, quero a nova colunapurch_revenue para mostrar a receita gerada pela compra em todas as linhas. Como solução alternativa, também tentei introduzir um identificador de comprapurch_id que é incrementado cada vez que uma compra foi feita. Portanto, isso é listado apenas como referência.

+-------+--------+-------+-------------+--------+
|user_id|visit_id|revenue|purch_revenue|purch_id|
+-------+--------+-------+-------------+--------+
|      1|       1|      0|          100|       1|
|      1|       2|      0|          100|       1|
|      1|       3|      0|          100|       1|
|      1|       4|    100|          100|       1|
|      1|       5|      0|          100|       2|
|      1|       6|      0|          100|       2|
|      1|       7|    200|          100|       2|
|      1|       8|      0|          100|       3|
|      1|       9|     10|          100|       3|
+-------+--------+-------+-------------+--------+

Eu tentei usar olag/lead funcionar assim:

user_timeline = Window.partitionBy("user_id").orderBy("visit_id")
find_rev = fn.when(fn.col("revenue") > 0,fn.col("revenue"))\ 
  .otherwise(fn.lead(fn.col("revenue"), 1).over(user_timeline))
df.withColumn("purch_revenue", find_rev)

Isso duplica a coluna de receita serevenue > 0 e também puxa-o por uma linha. Claramente, posso encadear isso para um N finito, mas isso não é uma solução.

Existe uma maneira de aplicar isso recursivamente atérevenue > 0?Como alternativa, existe uma maneira de incrementar um valor com base em uma condição? Eu tentei descobrir uma maneira de fazer isso, mas lutei para encontrar uma.