Pyspark有条件运行的窗口

Question

我有一个看起来像这样的数据框，

starttime           |  endtime            |
2022-01-01 03:25:53 | 2022-01-01 03:25:52 |
2022-01-01 03:25:53 | 2022-01-01 03:25:52 |
2022-01-01 03:25:53 | 2022-01-01 03:25:52 |
2022-01-01 03:25:55 | 2022-01-01 03:25:54 |
2022-01-01 03:25:57 | 2022-01-01 03:25:57 |

我想添加另一列是ROWNUMBER/等级，如果开始时间和末日之间的差异大于1秒，则会增加。

因此，预期的结果应该看起来像这样

starttime           |  endtime            | group |
2022-01-01 03:25:53 | 2022-01-01 03:25:52 | 1     |
2022-01-01 03:25:53 | 2022-01-01 03:25:52 | 1     |
2022-01-01 03:25:53 | 2022-01-01 03:25:52 | 1     |
2022-01-01 03:25:55 | 2022-01-01 03:25:54 | 2     |
2022-01-01 03:25:57 | 2022-01-01 03:25:57 | 3     |

Answer 1

据我了解，条目将按启动时间订购，而当启动时间和下一个终端时间之间的差异大于第二个时，我们需要组+1。

因此，根据上述要求，我们可以使用滞后和运行总和实现相同

test=spark.createDataFrame([("2022-01-01 03:25:53","2022-01-01 03:25:52"),("2022-01-01 03:25:53","2022-01-01 03:25:52"),("2022-01-01 03:25:53","2022-01-01 03:25:52"),("2022-01-01 03:25:55","2022-01-01 03:25:54"),("2022-01-01 03:25:57","2022-01-01 03:25:57")],"starttime string,endtime string")

test.withColumn("starttime",to_timestamp("starttime")).withColumn("endtime",to_timestamp("starttime")).\
withColumn("diff",when(unix_timestamp("endtime")-unix_timestamp(lag("starttime",1).over(Window.orderBy("starttime")))>1,1).otherwise(0)).\
withColumn("group",1+sum("diff").over(Window.orderBy("starttime"))).\
drop("diff").show()

#output
+-------------------+-------------------+-----+
|          starttime|            endtime|group|
+-------------------+-------------------+-----+
|2022-01-01 03:25:53|2022-01-01 03:25:53|    1|
|2022-01-01 03:25:53|2022-01-01 03:25:53|    1|
|2022-01-01 03:25:53|2022-01-01 03:25:53|    1|
|2022-01-01 03:25:55|2022-01-01 03:25:55|    2|
|2022-01-01 03:25:57|2022-01-01 03:25:57|    3|
+-------------------+-------------------+-----+