如何获得两个 Pyspark 日期之间的简单月份差异？ （与 SAS intck 的方式相同）

Question

我需要找出 Pyspark 中两个日期之间的差异 - 但模仿 SAS

intck

函数的行为。

我在下面列出了差异。

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)

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 和 4 时，SAS 计算两个日期之间的月份与 pyspark 不同。区别在于该月的某天之前和之后。

我想出了这个技巧：使用当天所在的 Spark 月 <= 24, otherwise add one.

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|
+----------+---------------+---------+

还有其他更好的解决方案吗？

Answer 1

我们可以获取年份差异、月份差异，然后将 intck 模拟为：

year_diff * 12 + month_diff

.

尝试以下方法，看看它是否适合您：

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)

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 = df.withColumn('year_diff', F.year(F.lit(ref_date)) - F.year(F.col('dt')))
df = df.withColumn('month_diff', F.month(F.lit(ref_date)) - F.month(F.col('dt')))
df = df.withColumn("new_spark_months_diff", F.expr("year_diff * 12 + month_diff"))
df.show()