PySpark中计算日期数组的相邻日期天数差值

需求说明 #

给定PySpark DataFrame中一列dates（类型为array<timestamp>，元素为有序日期），需要新增date_diffs列（类型为array<integer>），存储数组中相邻日期的天数差值。

实现代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, zip_with, datediff, expr

# 初始化Spark会话
spark = SparkSession.builder.appName("DateArrayDiff").getOrCreate()

# 构造示例数据
sample_data = [
    (["2012-02-13T00:00:00.000+0000", "2015-01-29T00:00:00.000+0000", "2018-07-29T00:00:00.000+0000", "2020-07-29T00:00:00.000+0000"],)
]

# 创建DataFrame
df = spark.createDataFrame(sample_data, ["dates"])

# 新增日期差值列
df = df.withColumn(
    "date_diffs",
    zip_with(
        col("dates"),
        expr("slice(dates, 2, size(dates) - 1)"),
        lambda prev_date, curr_date: datediff(curr_date.cast("date"), prev_date.cast("date"))
    )
)

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

关键逻辑说明 #

• zip_with函数：将原数组与它的子数组（去掉第一个元素）按位置配对，每一组对应相邻的两个日期
• slice函数：从原数组的第2个元素开始（PySpark数组索引从1起始），截取长度为原数组长度减1的子数组，确保和原数组配对后得到相邻元素对
• datediff函数：将时间戳转为date类型后计算天数差，避免时区差异导致的计算误差，返回整数结果

输出结果 #

+----------------------------------------------------------------------------------------------------+----------------+
|dates                                                                                               |date_diffs      |
+----------------------------------------------------------------------------------------------------+----------------+
|[2012-02-13 00:00:00, 2015-01-29 00:00:00, 2018-07-29 00:00:00, 2020-07-29 00:00:00]|[1081, 1277, 731]|
+----------------------------------------------------------------------------------------------------+----------------+

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