Разница во времени Pyspark, основанная на значениях столбцов

Question

У меня есть фрейм данных pyspark с четырьмя столбцами отметок времени (clock_in, clock_out, sign_in и sign_out).Мне нужен третий столбец (duration), который представляет собой разницу между двумя столбцами, но для которых я использую разницу, зависит от значений самих столбцов в соответствии со следующими правилами:

1. Если clock_in > clock_out, то duration = clock_in - clock_out
2. Остальное, если sign_in > clock_out, то duration = sign_out - clock_out
3. Остальное, если clock_in > sign_in, то duration = clock_in - clock_out
4. В противном случае, duration = 0

Я попытался определить UDF и применить его следующим образом:

from pyspark.sql import *

def calc_dur(clock_in, clock_out, sign_in, sign_out):
    if clock_in > clock_out:
        return clock_in - clock_out
    elif sign_out > clock_out:
        return sign_out - clock_out
    elif sign_in < clock_in:
        return clock_in - sign_in
    else:
        return 0

dur_udf = udf(calc_dur)

df2 = df.withColumn("duration", dur_udf(df.clock_in, df.clock_out, df.sign_in, df.sign_out))

Однако это заполняет мой столбец duration Timedeltas (например,, Timedelta: 0 days, 1740 seconds, 0 microseconds (total: 1740 seconds)), и когда я пытаюсь отфильтровать только строки с ненулевыми значениями в столбце duration, он возвращает ноль строк.

Как правильно это сделать?

Answer 1

Без использования UDF (что, на мой взгляд, медленнее), вы можете сделать это с помощью строгого синтаксиса Spark SQL, я считаю:

Столбцы должны быть заброшены заранее, если они имеют формат Timestamp. Таким образом, вы можете добавить этот фрагмент кода после каждого F.col ():

F.col("...").cast(LongType())

Ссылка на хороший SO ответ: https://stackoverflow.com/a/37058395/6434448

Этот код будет работать тогда, я надеюсь:

from pyspark.sql import functions as F
from pyspark.sql.types import LongType, TimestampType

df.withColumn("duration",
              F.when(F.col("clock_in") > F.col("clock_out"), F.col("clock_in") - F.col("clock_out")).otherwise(
              F.when(F.col("sign_out") > F.col("clock_out"), F.col("sign_out") - F.col("clock_out")).otherwise(
              F.when(F.col("sign_in") < F.col("clock_in"), F.col("clock_in") - F.col("sign_in")).otherwise(F.lit(0))
     ))
)