如何在PySpark中实现类似SAS intck函数的日期月份差计算？

PySpark的months_between函数会根据实际天数比例计算月份差值，而SAS的intck('month', start, end)函数是统计两个日期之间跨越的月份边界数量，二者逻辑差异会导致结果不一致。

测试数据与结果对比 #

import pyspark.sql.functions as F
import datetime

ref_date = '2023-02-24'

Data = [
    (1, datetime.date(2023, 1, 23), 1),
    (2, datetime.date(2023, 1, 24), 1),
    (3, datetime.date(2023, 1, 30), 1),
    (4, datetime.date(2022, 11, 30), 3),
    (5, datetime.date(2022, 11, 11), 3)
]

col = ['id', 'dt', 'SAS_months_diff']

df = spark.createDataFrame(Data, col)

# PySpark原生计算
df = df.withColumn('spark_months_diff', F.months_between(F.lit(ref_date), F.col('dt')).cast('integer'))
df = df.withColumn('spark_months_diff_x', F.months_between(F.lit(ref_date), F.col('dt')))

df.show()

输出结果：

+---+----------+---------------+-----------------+-------------------+
| id|        dt|SAS_months_diff|spark_months_diff|spark_months_diff_x|
+---+----------+---------------+-----------------+-------------------+
|  1|2023-01-23|              1|                1|         1.03225806|
|  2|2023-01-24|              1|                1|                1.0|
|  3|2023-01-30|              1|                0|         0.80645161|
|  4|2022-11-30|              3|                2|         2.80645161|
|  5|2022-11-11|              3|                3|         3.41935484|
+---+----------+---------------+-----------------+-------------------+

可见id=3和id=4的记录中，PySpark结果与SAS结果存在差异，核心原因是months_between按实际天数比例计算，而intck统计月份边界跨越数。

临时解决方案 #

你提到的临时方案通过判断日期的日部分是否大于参考日的日部分，对months_between的结果进行调整：

day = ref_date.split('-')[2]
df = df.withColumn('new_month', 
                   F.when(F.dayofmonth('dt') <= day, 
                          F.months_between(F.lit(ref_date), F.col('dt')).cast('integer'))
                   .otherwise(F.months_between(F.lit(ref_date), F.col('dt')).cast('integer') + 1))

df.select('dt', 'SAS_months_diff', 'new_month').show()

输出结果：

+----------+---------------+---------+
|        dt|SAS_months_diff|new_month|
+----------+---------------+---------+
|2023-01-23|              1|        1|
|2023-01-24|              1|        1|
|2023-01-30|              1|        1|
|2022-11-30|              3|        3|
|2022-11-11|              3|        3|
+----------+---------------+---------+

更优解决方案 #

直接通过提取日期的年、月部分计算差值，完全匹配SAS intck('month', dt, ref_date)的逻辑，无需依赖months_between的浮点计算，效率更高且逻辑更直观：

ref_date_lit = F.lit(ref_date)
# 计算参考日期的年月总数值
ref_year_month = F.year(ref_date_lit) * 12 + F.month(ref_date_lit)
# 计算目标日期的年月总数值
dt_year_month = F.year(F.col('dt')) * 12 + F.month(F.col('dt'))
# 差值即为跨越的月份边界数
df = df.withColumn('intck_month_diff', ref_year_month - dt_year_month)

df.select('dt', 'SAS_months_diff', 'intck_month_diff').show()

输出结果：

+----------+---------------+----------------+
|        dt|SAS_months_diff|intck_month_diff|
+----------+---------------+----------------+
|2023-01-23|              1|               1|
|2023-01-24|              1|               1|
|2023-01-30|              1|               1|
|2022-11-30|              3|               3|
|2022-11-11|              3|               3|
+----------+---------------+----------------+

这种方法直接通过整数运算得到结果，避免了浮点转换可能带来的精度问题，同时计算逻辑完全贴合SAS intck的边界统计规则。

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