Spark 2.4中如何获取指定月份范围内的周起止日期

需求 #

找出指定月份内的周起始日期和结束日期，结果不得超出给定月份的范围。

输入DataFrame #

df = spark.createDataFrame([('1', '2022-07-01'), ('2', '2022-07-30'), ('3', '2022-07-31')], ["id", "date"])
df.createOrReplaceTempView("date_table")

# 数据预览：
# +---+----------+
# | id|      date|
# +---+----------+
# | 1 |2022-07-01|
# | 2 |2022-07-30|
# | 3 |2022-07-31|
# +---+----------+

预期输出 #

+---+----------+-----------+-----------+
| id|      date|week_start | week_end  |
+---+----------+-----------+-----------+
| 1 |2022-07-01|2022-07-01 |2022-07-02 |
| 2 |2022-07-30|2022-07-24 |2022-07-30 |
| 3 |2022-07-31|2022-07-31 |2022-07-31 |
+---+----------+-----------+-----------+

现有问题 #

原SQL查询未做月份边界限制，会出现周起止日期超出目标月份的情况：

spark.sql(""" select id,date,date_sub(date,dayofweek-1) as week_start,
                date_add(date, 7-dayofweek) as week_end FROM
                (select id,date,dayofweek(date) as dayofweek
                FROM date_table) T""")

比如id=1的记录，计算出的周起始会落到2022年6月；id=3的记录，周结束会落到2022年8月，不符合需求。

解决方案 #

通过trunc获取当月第一天、last_day获取当月最后一天，再用GREATEST和LEAST函数限制周起止日期的边界，确保结果不会超出目标月份：

SELECT 
    id,
    date,
    GREATEST(date_sub(date, dayofweek(date) - 1), trunc(date, 'month')) AS week_start,
    LEAST(date_add(date, 7 - dayofweek(date)), last_day(date)) AS week_end
FROM date_table

逻辑说明 #

• trunc(date, 'month')：获取当前日期所在月份的第一天
• last_day(date)：获取当前日期所在月份的最后一天
• GREATEST(计算出的周起始, 当月第一天)：保证周起始不会早于当月第一天
• LEAST(计算出的周结束, 当月最后一天)：保证周结束不会晚于当月最后一天

执行该查询后，结果将完全符合预期输出。

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