如何在PySpark中修复数据集的日期重叠并生成连续日期

问题描述 #

我有一个包含Key1、Key2、Key3、Key4、Value、Date1和Date2字段的数据集，其中Date1与Date2存在日期重叠问题。原始数据如下：

+------+------+------+------+-------+------------+------------+
| Key1 | Key2 | Key3 | Key4 | Value |   Date1    |   Date2    |
+------+------+------+------+-------+------------+------------+
| k1   | k2   | k3   | k4   |    10 | 2022-01-01 | 2026-01-30 |
| k1   | k2   | k3   | k4   |    12 | 2022-06-05 | 2026-01-10 |
| k1   | k2   | k3   | k4   |    14 | 2022-08-07 | 2026-01-15 |
+------+------+------+------+-------+------------+------------+

需要修复日期重叠，使日期范围连续，调整规则为：new_date2 = 下一条记录的old_date1 - 1，期望结果如下：

+------+------+------+------+-------+------------+------------+
| Key1 | Key2 | Key3 | Key4 | Value |   Date1    |   Date2    |
+------+------+------+------+-------+------------+------------+
| k1   | k2   | k3   | k4   |    10 | 2022-01-01 | 2022-06-04 |
| k1   | k2   | k3   | k4   |    12 | 2022-06-05 | 2022-08-06 |
| k1   | k2   | k3   | k4   |    14 | 2022-08-07 | 2026-01-15 |
+------+------+------+------+-------+------------+------------+

解决方案 #

使用SQL的窗口函数LEAD可以高效实现这个需求，以下是示例代码：

WITH ranked_data AS (
    SELECT 
        Key1, Key2, Key3, Key4, Value, Date1, Date2,
        -- 按Key分组、Date1升序，获取下一条记录的Date1
        LEAD(Date1) OVER (PARTITION BY Key1, Key2, Key3, Key4 ORDER BY Date1) AS next_date1
    FROM your_table
)
SELECT 
    Key1, Key2, Key3, Key4, Value,
    Date1,
    -- 调整Date2：有下一条记录则取其Date1减1，否则保留同组最大Date2
    CASE 
        WHEN next_date1 IS NOT NULL THEN DATEADD(day, -1, next_date1)
        ELSE (SELECT MAX(Date2) FROM your_table WHERE Key1 = rd.Key1 AND Key2 = rd.Key2 AND Key3 = rd.Key3 AND Key4 = rd.Key4)
    END AS Date2
FROM ranked_data rd
ORDER BY Date1;

代码说明 #

1. LEAD窗口函数：用于匹配同组（Key1-Key4完全一致）且按Date1排序的下一条记录的起始日期。
2. CASE逻辑处理：
   • 若存在下一条记录，将当前记录的Date2设置为下一条记录的Date1减1天，确保日期区间连续无重叠。
   • 若为组内最后一条记录，取同组所有记录中最大的Date2作为最终结束日期，对应示例中最后一条的2026-01-15。

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