Как рассчитать время до следующего сбоя

Question

Моя цель - вычислить для каждой временной отметки время до следующего отказа.

Я реализовал метод, чтобы сделать это, сначала создав рамку даты со всеми временными метками сбоев и присоединив ее к условию

H.show()

            Timestamp      state
0   2017-03-30 03:51:57.900 OK
1   2017-03-30 03:51:58.000 OK
2   2017-03-30 03:51:58.100 NOK
3   2017-03-30 03:51:58.200 OK
4   2017-03-30 03:51:58.300 NOK

failures = H.filter(H['state'] == 'NOK').select('Timestamp')\
.selectExpr('Timestamp as Failure_time')
M = H.join(failures, F.col('Timestamp') <= F.col("Failure_time"))
               Timestamp    state   Failure_time
0   2017-03-30 03:51:57.900 OK  2017-03-30 03:51:58.100
1   2017-03-30 03:51:57.900 OK  2017-03-30 03:51:58.300
2   2017-03-30 03:51:58.000 OK  2017-03-30 03:51:58.100
3   2017-03-30 03:51:58.000 OK  2017-03-30 03:51:58.300
4   2017-03-30 03:51:58.100 NOK 2017-03-30 03:51:58.100
5   2017-03-30 03:51:58.100 NOK 2017-03-30 03:51:58.300
6   2017-03-30 03:51:58.200 OK  2017-03-30 03:51:58.300
7   2017-03-30 03:51:58.300 NOK 2017-03-30 03:51:58.300

M= M.groupBy('Timestamp').agg(F.min('Failure_time').alias('Next_Failure'))\
.orderBy('Timestamp')

               Timestamp           Next_Failure
0   2017-03-30 03:51:57.900 2017-03-30 03:51:58.100
1   2017-03-30 03:51:58.000 2017-03-30 03:51:58.100
2   2017-03-30 03:51:58.100 2017-03-30 03:51:58.100
3   2017-03-30 03:51:58.200 2017-03-30 03:51:58.300
4   2017-03-30 03:51:58.300 2017-03-30 03:51:58.300

Этот метод выполняет желаемую работу, но занимает очень много времени для информационного кадра с только 1,5 миллионами записей (что-то вроде двух дней с очень мощной машиной).

Может ли кто-нибудь помочь с более эффективным методом, именно я хотел бы избежать полного объединения и агрегирования groupBy.

Большое спасибо