Как сканировать столбец, чтобы получить новый столбец в Pyspark DataFrame

Question

У меня есть Pyspark DataFrame с двумя столбцами: sendtime и charge_state, если обвинение_state меняется с «выключено» на «включено», начинается новый цикл зарядки.

Теперь я хочу пометить каждый цикл зарядки какдай мне вывод.

Ввод:

+-------------------+------------+
|           sendtime|charge_state|
+-------------------+------------+
|2018-03-02 08:00:00|          on|
...
|2018-03-02 09:42:32|          on|
|2018-03-02 09:42:33|          on|
|2018-03-02 09:42:34|          on|
|2018-03-02 09:42:35|         off|
|2018-03-02 09:42:36|         off|
...
|2018-03-02 10:11:12|         off|
|2018-03-02 10:11:13|          on|
|2018-03-02 10:11:14|          on|
...

Выход:

+-------------------+------------+---------------+
|           sendtime|charge_state|charge_cycle_ID|
+-------------------+------------+---------------+
|2018-03-02 08:00:00|          on|             c1|
...
|2018-03-02 09:42:32|          on|             c1|
|2018-03-02 09:42:33|          on|             c1|
|2018-03-02 09:42:34|          on|             c1|
|2018-03-02 09:42:35|         off|             c1|
|2018-03-02 09:42:36|         off|             c1|
...
|2018-03-02 10:11:12|         off|             c1|
|2018-03-02 10:11:13|          on|             c2|
|2018-03-02 10:11:14|          on|             c2|
...

Answer 1

Вы можете использовать функции Window для выполнения этой задачи:

from pyspark.sql import functions as F
from pyspark.sql import Window

df.withColumn(
    'charge_state_lag', 
    F.lag('charge_state').over(Window.partitionBy().orderBy('sendtime'))
).withColumn(
    'fg', 
    F.when((F.col("charge_state")=="on")&(F.col("charge_state_lag")=="off"),1).otherwise(0)
).select(
    'sendtime',
    'charge_state',
    F.concat(
        F.lit('C'),
        (F.sum('fg').over(Window.partitionBy().orderBy('sendtime'))+1).cast('string')
    ).alias("charge_cycle_ID")
).show()