基于列条件的 Pyspark 计数器

Question

您好，我想根据 Tag1 列的条件创建一个带有计数器的新列。

我有这个：

Time	Tag1
1	0
2	1
3	1
4	1
5	0
6	0
7	1
8	1
9	1
10	1
11	0
12	0

我想要这样：

Time	Tag1	Counter
1	0	0
2	1	1
3	1	2
4	1	3
5	0	0
6	0	0
7	1	1
8	1	2
9	1	3
10	1	4
11	0	0
12	0	0

我尝试使用 function.when(df.Tag1 == 1, function.lag(df.Tag1)+1).otherwise(0) 但不起作用。

有什么想法吗？

多谢

Answer 1

窗函数

    new = (df.withColumn('Counter',(col('Tag1')=='0'))#Create Bool
       .withColumn('Counter', F.sum(F.col('Counter').cast('integer')).over(Window.partitionBy().orderBy().rowsBetween(-sys.maxsize, 0)))# Create Group by summing bool
.withColumn('Counter', when(col('Tag1')==0, col('Tag1')).otherwise(F.sum('Tag1').over(Window.partitionBy('Counter').orderBy().rowsBetween(-sys.maxsize, 0))))#Conditionally add
       
      )

new.show()

+----+----+-------+
|Time|Tag1|Counter|
+----+----+-------+
|   1|   0|      0|
|   2|   1|      1|
|   3|   1|      2|
|   4|   1|      3|
|   5|   0|      0|
|   6|   0|      0|
|   7|   1|      1|
|   8|   1|      2|
|   9|   1|      3|
|  10|   1|      4|
|  11|   0|      0|
|  12|   0|      0|
+----+----+-------+