Create unique IDs from a run of 0s and 1s

consecutive_group_ids.py

def consecutive_group_ids(df, id_col: str, time_col: str, value_col: str, event_id_col: str="EventID"):
    w1 = Window.partitionBy(id_col).orderBy(time_col)
    wcumsum = (
        Window
        .partitionBy(id_col)
        .orderBy(time_col)
        .rangeBetween(Window.unboundedPreceding, 0)
    )

    res = (
        df
        .orderBy([id_col, time_col])
        .withColumn('Diff', F.abs(F.col(value_col) - F.lag(F.col(value_col)).over(w1)))
        .fillna({"Diff":1})
        .withColumn('EventGroup', F.sum('Diff').over(wcumsum))
        .drop('Diff')
        .withColumn(event_id_col, F.dense_rank().over(Window.orderBy([id_col,"EventGroup"])))
    )
    return res