基于列值的 Pyspark 时间戳差异

Question

我有一个带有四个时间戳列（clock_in、clock_out、sign_in 和 sign_out）的 pyspark 数据框。我需要第三列 (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))

但是，这会用 Timedeltas（例如，Timedelta: 0 days, 1740 seconds, 0 microseconds (total: 1740 seconds)）填充我的 duration 列，并且当我尝试仅过滤到 duration 列中具有非零值的行时，它会返回零行。

这样做的正确方法是什么？

Answer 1

不必使用 UDF（我认为它更慢），我相信您可以使用严格的 Spark SQL 语法来做到这一点：

列应预先转换（如果它们是时间戳格式）。所以你可以在每个 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))
     ))
)