PySpark pandas中两日期列天数差的高效计算方法

问题背景 #

在原生pandas中，将两列datetime64类型的日期相减后，可通过.dt.days直接获取天数差；但在PySpark pandas（以下简称ps）中，两日期列直接相减会返回以秒数为单位的整数（如示例中的[-259200, 0, -86400, -86400, -86400, -86400]），无法直接得到天数格式。尝试将该整数列转为datetime类型会得到1970年左右的无效日期，而将数据转回原生pandas处理虽能得到正确结果，但大数据量下会带来额外开销、延长运行时间，因此需要高效的分布式实现方式。

示例代码：

import pandas as pd
import pyspark.pandas as ps

data = {
    "d1": [
        "2019-05-18",
        "2019-06-21",
        "2019-05-08",
        "2019-05-22",
        "2019-11-20",
        "2019-05-29",
    ],
    "d2": [
        "2019-05-21",
        "2019-06-21",
        "2019-05-09",
        "2019-05-23",
        "2019-11-21",
        "2019-05-30",
    ],
}
df = pd.DataFrame(data)
df[["d1", "d2"]] = df[["d1", "d2"]].astype("datetime64")
pdf = ps.from_pandas(df)

# 原生pandas的正确结果
df["diff"] = (df["d1"] - df["d2"]).dt.days
# PySpark pandas直接相减得到秒数
pdf["diff_seconds"] = pdf["d1"] - pdf["d2"]

原生pandas输出的天数差：[-3, 0, -1, -1, -1, -1]

高效解决方案 #

方法1：秒数直接转换为天数 #

既然相减结果是秒数，直接除以一天的秒数（86400）并做整数除法即可，简洁且无额外依赖：

pdf["diff_days"] = (pdf["d1"] - pdf["d2"]) // 86400

该方法保证结果为整数天数，与原生pandas输出完全一致。

方法2：使用PySpark内置函数datediff #

通过PySpark pandas的ps.sql接口调用Spark原生的datediff函数，分布式执行更适合大数据量场景：

pdf["diff_days"] = ps.sql("SELECT datediff(d1, d2) FROM {pdf}")

注意：datediff参数顺序为datediff(endDate, startDate)，结果为endDate - startDate的天数，与示例中d1 - d2逻辑匹配。

方法3：利用PySpark pandas的dt属性 #

将日期相减结果转为Timedelta类型，再通过.dt.days获取天数，贴近原生pandas使用习惯：

pdf["diff_days"] = ps.Series(pdf["d1"] - pdf["d2"], dtype="timedelta64[ns]").dt.days

验证结果 #

以上三种方法最终得到的diff_days列结果均为：[-3, 0, -1, -1, -1, -1]，与原生pandas输出一致，且无需将数据拉回本地处理。

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