Spark DataFrame两日期相对周数计算的优化方案问询

我有一个包含date1和date2两列日期的Spark DataFrame，需要计算两个日期之间的相对周数（结果需加1周）。其中date2始终早于date1，日期最大差值不超过1年，若date2属于上一年，需在结果中加52。我已实现需求但方案较为繁琐，现有代码如下：

from pyspark.sql.functions import date_format, udf
from pyspark.sql.types import IntegerType

def fixWeeks(w):
    if w < 0:
        w = w + 52
    return w

week_filter = udf(fixWeeks,IntegerType())
df_week = df.withColumn("TmpWeek", date_format("date1",'w').cast('int')-date_format("date2",'w').cast('int')+1)
df_week = df_week.withColumn('week',week_filter(df_week.TmpWeek)).drop('TmpWeek')

请问是否有更优的实现方式？如何用单个函数完成该需求？

更优实现方案：避免UDF，使用Spark内置函数 #

你的需求完全可以通过Spark的内置函数组合实现，不需要自定义UDF——这不仅能让代码更简洁，还能避免Python UDF带来的性能开销（毕竟Spark内置函数是在JVM端执行，分布式场景下效率远高于Python UDF）。

核心思路

你的逻辑本质是：计算date1周数 - date2周数 + 1，如果结果为负（说明date2在上一年）则加52。利用模运算的特性，我们可以用一行代码完成这个逻辑：

from pyspark.sql.functions import weekofyear, expr

# 直接计算目标周数，无需临时列和UDF
df_week = df.withColumn(
    "week",
    expr("(weekofyear(date1) - weekofyear(date2) + 1 + 52) % 52")
)

为什么这个方案可行？

• 当date1和date2在同一年时，weekofyear(date1) - weekofyear(date2) + 1结果为正，加52再取模52后仍等于原结果，符合需求。
• 当date2在上一年时，比如date1是当年第1周，date2是上年第52周，计算得1 - 52 +1 = -50，加52后为2，模52还是2，刚好满足“加52”的要求。
• 由于你明确说明日期最大差值不超过1年，date2始终早于date1，所以这个公式不会出现超出预期的结果。

如果你想保持和原代码一致的周数计算方式

原代码用date_format('w')获取周数，如果你担心weekofyear的行为和它有差异（注：两者默认都是ISO周数，大部分场景下一致），可以替换成date_format的写法：

from pyspark.sql.functions import expr

df_week = df.withColumn(
    "week",
    expr("(date_format(date1, 'w')::int - date_format(date2, 'w')::int + 1 + 52) % 52")
)

不用expr的函数链式写法

如果你更喜欢纯函数调用的方式，也可以这样写：

from pyspark.sql.functions import weekofyear, mod, lit

df_week = df.withColumn(
    "week",
    mod(
        weekofyear("date1") - weekofyear("date2") + lit(1) + lit(52),
        lit(52)
    )
)

这些方案都不需要临时列和UDF，代码更简洁高效，完全满足你的需求。

内容的提问来源于stack exchange，提问作者Eyal S.