PySpark Window function to check based on condition a lag row - pyspark

Question

I would need to take (if present and last) STAT 200.
I would like to take its date only if last status and for this case I know how to do it, it's ok.
But i need first to check it (if it not the last status) and then,
if it were not the last, I would like to consider the date of the last state as 200 but only if the next state is not minor (for example 150 as in the next case). In this case I wouldn't want to consider it.
How to do in PySpark? Thank you very much
ID
STAT
TIME
1
100
17/11/2021
1
200
18/11/2021
1
150
19/11/2021
1
100
20/11/2021
2
200
20/11/2021
In this case I don't want to consider ID 1 (status 150 < 200). So I take only ID 2.
ID
STAT
TIME
1
100
17/11/2021
1
200
18/11/2021
1
350
19/11/2021
1
400
20/11/2021
2
200
20/11/2021
In this case I want to consider status 200 and its details (status 350 > 200 for ID 1; just 200 for ID 2)
Thank u for support!

Answer 1

You can specify a default value to the lead function and then handle last row with STAT = 200 and non-last row with STAT=200 using the same logic.
from pyspark.sql import functions as F
from pyspark.sql import Window
window_spec = Window.partitionBy("ID").orderBy("TIME")
df.withColumn("LEAD_STAT", F.lead("STAT", default=201).over(window_spec)).filter((F.col("STAT") == 200) & (F.col("LEAD_STAT") > 200)).show()
Example
Scneario One
from datetime import datetime
from pyspark.sql import functions as F
from pyspark.sql import Window
data = [(1, 100, datetime.strptime("17/11/2021", "%d/%m/%Y"),),
(1, 200, datetime.strptime("18/11/2021", "%d/%m/%Y"),),
(1, 150, datetime.strptime("19/11/2021", "%d/%m/%Y"),),
(1, 100, datetime.strptime("20/11/2021", "%d/%m/%Y"),),
(2, 200, datetime.strptime("20/11/2021", "%d/%m/%Y"),),]
df = spark.createDataFrame(data, ("ID", "STAT", "TIME"))
window_spec = Window.partitionBy("ID").orderBy("TIME")
df.withColumn("LEAD_STAT", F.lead("STAT", default=201).over(window_spec)).filter((F.col("STAT") == 200) & (F.col("LEAD_STAT") > 200)).show()
Output
+---+----+-------------------+---------+
| ID|STAT| TIME|LEAD_STAT|
+---+----+-------------------+---------+
| 2| 200|2021-11-20 00:00:00| 201|
+---+----+-------------------+---------+
Scneario Two
from datetime import datetime
from pyspark.sql import functions as F
from pyspark.sql import Window
data = [(1, 100, datetime.strptime("17/11/2021", "%d/%m/%Y"),),
(1, 200, datetime.strptime("18/11/2021", "%d/%m/%Y"),),
(1, 350, datetime.strptime("19/11/2021", "%d/%m/%Y"),),
(1, 400, datetime.strptime("20/11/2021", "%d/%m/%Y"),),
(2, 200, datetime.strptime("20/11/2021", "%d/%m/%Y"),),
]
df = spark.createDataFrame(data, ("ID", "STAT", "TIME"))
window_spec = Window.partitionBy("ID").orderBy("TIME")
df.withColumn("LEAD_STAT", F.lead("STAT", default=201).over(window_spec)).filter((F.col("STAT") == 200) & (F.col("LEAD_STAT") > 200)).show()
Output
+---+----+-------------------+---------+
| ID|STAT| TIME|LEAD_STAT|
+---+----+-------------------+---------+
| 1| 200|2021-11-18 00:00:00| 350|
| 2| 200|2021-11-20 00:00:00| 201|
+---+----+-------------------+---------+