PySpark 3.x计算DataFrame嵌套时间戳数组的秒级时间差

需求说明 #

给定一个包含嵌套时间戳数组的PySpark DataFrame（Spark 3.x版本），需计算每个子数组内两个时间戳的秒级差值，生成time_diff列。

示例输入 #

+------------------------------------------------------------------------------------------------------------------------------------+
|timestamp                                                                                                                           |
+------------------------------------------------------------------------------------------------------------------------------------+
|[[2022-01-01 12:00:00, 2022-01-02 15:30:00]]                                                                                        |
|[[2022-01-01 12:00:00, 2022-01-02 14:30:00], [2022-01-02 12:00:00, 2022-01-03 19:30:00], [2022-01-02 12:00:00, 2022-01-03 15:30:00]]|
|[[2022-01-01 12:00:00, 2022-01-02 16:30:00], [2022-01-03 12:00:00, 2022-01-04 17:30:00]]                                            |
|[]                                                                                                                                  |
+------------------------------------------------------------------------------------------------------------------------------------+

期望输出 #

+------------------------------------------------------------------------------------------------------------------------------------+------------------------------+
|timestamp                                                                                                                           |    time_diff                 |
+------------------------------------------------------------------------------------------------------------------------------------+------------------------------+
|[[2022-01-01 12:00:00, 2022-01-02 15:30:00]]                                                                                        |[99000.0]                     |
|[[2022-01-01 12:00:00, 2022-01-02 14:30:00], [2022-01-02 12:00:00, 2022-01-03 19:30:00], [2022-01-02 12:00:00, 2022-01-03 15:30:00]]|[95400.0,199800.0,185400.0]   |
|[[2022-01-01 12:00:00, 2022-01-02 16:30:00], [2022-01-03 12:00:00, 2022-01-04 17:30:00]]                                            |[189000.0,192600.0]           |
|[]                                                                                                                                  |[]                            |
+------------------------------------------------------------------------------------------------------------------------------------+------------------------------+

解决方案 #

利用Spark 3.x支持的高阶函数transform，无需自定义UDF即可高效实现需求：

步骤1：导入依赖 #

from pyspark.sql import functions as F

步骤2：计算时间差 #

假设你的DataFrame名为df，执行以下代码生成time_diff列：

df = df.withColumn(
    "time_diff",
    F.transform(
        F.col("timestamp"),
        # 对每个子数组，计算第二个时间戳与第一个的秒级差值
        lambda sub_arr: F.unix_timestamp(sub_arr.getItem(1)) - F.unix_timestamp(sub_arr.getItem(0))
    ).cast("array<double>")  # 转换为示例中的浮点数组类型
)

补充：若原始时间戳为字符串类型 #

如果timestamp列的元素是字符串而非Spark timestamp类型，需先做类型转换：

# 将字符串数组转为timestamp数组
df = df.withColumn(
    "timestamp",
    F.transform(
        F.col("timestamp"),
        lambda sub_arr: F.array(
            F.to_timestamp(sub_arr.getItem(0)),
            F.to_timestamp(sub_arr.getItem(1))
        )
    )
)

完整测试代码 #

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

spark = SparkSession.builder.appName("TimestampDiff").getOrCreate()

# 构造示例数据
data = [
    ([[("2022-01-01 12:00:00"), ("2022-01-02 15:30:00")]],),
    ([[("2022-01-01 12:00:00"), ("2022-01-02 14:30:00")], [("2022-01-02 12:00:00"), ("2022-01-03 19:30:00")], [("2022-01-02 12:00:00"), ("2022-01-03 15:30:00")]],),
    ([[("2022-01-01 12:00:00"), ("2022-01-02 16:30:00")], [("2022-01-03 12:00:00"), ("2022-01-04 17:30:00")]],),
    ([],)
]

df = spark.createDataFrame(data, ["timestamp"])

# 字符串转timestamp类型
df = df.withColumn(
    "timestamp",
    F.transform(
        F.col("timestamp"),
        lambda sub_arr: F.array(F.to_timestamp(sub_arr.getItem(0)), F.to_timestamp(sub_arr.getItem(1)))
    )
)

# 计算时间差
df = df.withColumn(
    "time_diff",
    F.transform(
        F.col("timestamp"),
        lambda sub_arr: F.unix_timestamp(sub_arr.getItem(1)) - F.unix_timestamp(sub_arr.getItem(0))
    ).cast("array<double>")
)

# 查看结果
df.show(truncate=False)

内容的提问来源于stack exchange，提问作者pks_zon