如果火花中每个组满足特定条件，则将值添加到列中

Question

我想我有一些相当简单的东西。

我试图实现的是，如果满足某个条件，则按组给出一个增加数(排名？(。对于每组，它以1开头，若满足条件，则下一行为前一行的值+1。这种情况在组中越来越严重，每次满足条件时，加1。

下表可能更清楚地显示了这一点。(我试图创建的是"What_I_want"列(

group   to_add_number   what_i_want
aaaaaa  0                 1
aaaaaa  0                 1
aaaaaa  1                 2
aaaaaa  0                 2
aaaaaa  0                 2
aaaaaa  1                 3
aaaaaa  0                 3
aaaaaa  0                 3
bbbbbb  0                 1
bbbbbb  1                 2
bbbbbb  1                 3
bbbbbb  0                 3
cccccc  0                 1
cccccc  0                 1
cccccc  0                 1
cccccc  1                 2

我认为窗口函数(滞后(可能会做到这一点，但我做不到。

我尝试的是：

from pyspark.sql.functions import lit,when,lag,row_number
from pyspark.sql.window import Window
windowSpec=Window.partitionBy('group')
df=df.withColumn('tmp_rnk',lit(1))
df=df.withColumn('what_i_want',when(col('to_add_number')==0,lag('tmp_rnk').over(windowSpec)).otherwise(col('what_i_want')+1)
or
df=df.withColumn('tmp_rnk',lit(1))
df=df.withColumn('row_number_rank',row_number().over(windowSpec))
df=df.withColumn('what_i_want',when((col('to_add_number')==0)&(col('row_number_rank')==1)
,lit(1)
.when(col('to_add_number')==0)&(col('row_number_rank')>1),lag('what_i_want').over(windowSpec).otherwise(col('what_i_want')+1)

我尝试了几种变体，在stackoverflow上搜索"条件窗口函数"、"滞后、超前……"(，但什么都没用，或者我没有找到重复的问题。

Answer 1

若要获取列what_i_want，您可以在to_add_number上运行incremental sum并使用orderby列(order_id(。

from pyspark.sql import functions as F
from pyspark.sql.window import Window
df.withColumn("order_id", F.monotonically_increasing_id())
.withColumn("what_i_want", F.sum("to_add_number").over(Window().partitionBy("group").orderBy("order_id"))+1)
.orderBy("order_id").drop("order_id").show()

#+------+-------------+-----------+
#| group|to_add_number|what_i_want|
#+------+-------------+-----------+
#|aaaaaa|            0|          1|
#|aaaaaa|            0|          1|
#|aaaaaa|            1|          2|
#|aaaaaa|            0|          2|
#|aaaaaa|            0|          2|
#|aaaaaa|            1|          3|
#|aaaaaa|            0|          3|
#|aaaaaa|            0|          3|
#|bbbbbb|            0|          1|
#|bbbbbb|            1|          2|
#|bbbbbb|            1|          3|
#|bbbbbb|            0|          3|
#|cccccc|            0|          1|
#|cccccc|            0|          1|
#|cccccc|            0|          1|
#|cccccc|            1|          2|
#+------+-------------+-----------+