How to label rows in PySpark

I have the following dataframe in Pyspark:

I need the following output:

| ID | Timestamp | Event | Type |
| -------- | -------- |-----------
| 1 | 1657610298 | 0 | before
| 1 | 1657610299 | 0 | before
| 1 | 1657610300 | 0 | before
| 1 | 1657610301 | 1 | event
| 1 | 1657610302 | 0 | after
| 1 | 1657610303 | 0 | after
| 1 | 1657610304 | 0 | after
| 2 | 1657610298 | 0 | before
| 2 | 1657610299 | 0 | before
| 2 | 1657610300 | 0 | before
| 2 | 1657610301 | 1 | event
| 2 | 1657610302 | 0 | after
| 2 | 1657610303 | 0 | after
| 2 | 1657610304 | 0 | after

I tried the following function:

def label(df_):
    remove = ['type1']
    df_ = (
        df_
        .withColumn('type1', F.when((F.col("Event") == 0) & (F.lag(F.col("Event"), 1).over(Window.partitionBy('ID').orderBy('Timestamp')) == 1), 
                                         F.lit('after')))
        .withColumn('type2', F.when((F.col("isHypoProtectEnabled") == 0) & ((F.lag(F.col("Event"), 1).over(Window.partitionBy('ID').orderBy('Timestamp')) == 1) | (F.lag(F.col("type1"), 1).over(Window.partitionBy('ID').orderBy('Timestamp')) == 'after')), 
                                         F.lit('after')).otherwise(F.lit('before')))
    )
    df_ = df_.drop(*remove)
    return df_

What I get:

| **ID**   | **Timestamp** | **Event** | **type2** |
| -------- | --------      |-----------
| 1        | 1657610298    | 0         | before
| 1        | 1657610299    | 0         | before
| 1        | 1657610300    | 0         | before
| 1        | 1657610301    | 1         | event
| 1        | 1657610302    | 0         | after
| 1        | 1657610303    | 0         | after
| 1        | 1657610304    | 0         | *before <- error* should be after
| 2        | 1657610298    | 0         | before
| 2        | 1657610299    | 0         | before
| 2        | 1657610300    | 0         | before
| 2        | 1657610301    | 1         | event
| 2        | 1657610302    | 0         | after
| 2        | 1657610303    | 0         | after
| 2        | 1657610304    | 0         | *before <- error* should be after

Obviously is not working, because in order for it to label all the 'after' correctly I would have to loop the function... right now I'm only getting TWO 'after's per ID, when it should be more, I'm sure there's another way to do it, but I'm blocked please help me

Create temp_col based on the max values from unboundedPreceeding, current_row and then check

• if both columns are 0 then before
• if event column is 0 and if temp_col columns are 1 then after.
• else its an event.

Example:

df_windw = df.withColumn("temp_col",max(col("Event")).over(Window.partitionBy('ID').orderBy("timestamp").rowsBetween(Window.unboundedPreceding, Window.currentRow))).\
  withColumn("type2", when((col("event")== 0) & (col("temp_col")==0),lit("before")).\
    when((col("event")== 0) & (col("temp_col")==1),lit("after")).\
      otherwise(lit("event"))).\
        drop("temp_col")
df_windw.show(100,False)
#+---+----------+-----+------+
#|ID |timestamp |event|type2 |
#+---+----------+-----+------+
#|1  |1657610298|0    |before|
#|1  |1657610299|0    |before|
#|1  |1657610300|0    |before|
#|1  |1657610301|1    |event |
#|1  |1657610302|0    |after |
#|1  |1657610303|0    |after |
#|1  |1657610304|0    |after |
#|2  |1657610298|0    |before|
#|2  |1657610299|0    |before|
#|2  |1657610300|0    |before|
#|2  |1657610301|1    |event |
#|2  |1657610302|0    |after |
#|2  |1657610303|0    |after |
#|2  |1657610304|0    |after |
#+---+----------+-----+------+

Here is another solution, but without window function:

from pyspark.sql.functions import when

# Create sample data
data = [
    (1, 1657610298, 0),
    (1, 1657610299, 0),
    (1, 1657610300, 0),
    (1, 1657610301, 1),
    (1, 1657610302, 0),
    (1, 1657610303, 0),
    (1, 1657610304, 0),
    (2, 1657610298, 0),
    (2, 1657610299, 0),
    (2, 1657610300, 0),
    (2, 1657610301, 1),
    (2, 1657610302, 0),
    (2, 1657610303, 0),
    (2, 1657610304, 0)
]
# Create DataFrames from the sample data (all data and just events)
df_all = spark.createDataFrame(data, ["ID", "Timestamp", "Event"])
df_events = df_all.filter("Event==1")

# Join df_all with df_events based on ID
joined_df = df_all.join(
    df_events.withColumnRenamed('Timestamp', 'EventTimestamp').withColumnRenamed('Event', 'EventEvent'), "ID", "left")

# Add the "Type" column based on conditions
result_df = joined_df.withColumn(
    "Type",
    when(joined_df.Timestamp < joined_df.EventTimestamp, "before")
    .when(joined_df.Timestamp == joined_df.EventTimestamp, "event")
    .otherwise("after")
)

# Drop the duplicated columns created by join
result_df = result_df.drop("EventTimestamp").drop("EventEvent")

# Print the output
result_df.show(1000, False)

And the output is:

+---+----------+-----+------+                                                   
|ID |Timestamp |Event|Type  |
+---+----------+-----+------+
|1  |1657610298|0    |before|
|1  |1657610299|0    |before|
|1  |1657610300|0    |before|
|1  |1657610301|1    |event |
|1  |1657610302|0    |after |
|1  |1657610303|0    |after |
|1  |1657610304|0    |after |
|2  |1657610298|0    |before|
|2  |1657610299|0    |before|
|2  |1657610300|0    |before|
|2  |1657610301|1    |event |
|2  |1657610302|0    |after |
|2  |1657610303|0    |after |
|2  |1657610304|0    |after |
+---+----------+-----+------+