如何在PySpark中为分组生成3天滚动序列ID（无UDF）

解决方案代码（无UDF） #

针对你的需求，我们可以用递归CTE实现，完全兼容Databricks 13.1运行时，无需自定义UDF：

WITH RECURSIVE grouped_sorted AS (
    -- 对每个(phone_number, service)分组的日期排序，生成行号用于递归遍历
    SELECT 
        phone_number,
        service,
        event_date,
        ROW_NUMBER() OVER (PARTITION BY phone_number, service ORDER BY event_date) AS rn
    FROM test_data
),
sequence_generator AS (
    -- 初始化：每个分组的第一行，序列ID设为1，块起始日期为当前事件日期
    SELECT 
        phone_number,
        service,
        event_date,
        1 AS sequence_id,
        event_date AS block_start_date
    FROM grouped_sorted
    WHERE rn = 1
    
    UNION ALL
    
    -- 递归处理后续行：判断当前日期是否超出当前块的3天间隔，决定是否重置序列
    SELECT 
        g.phone_number,
        g.service,
        g.event_date,
        -- 若当前日期与块起始日期差≥3天，序列ID+1；否则沿用原序列ID
        CASE WHEN DATEDIFF(g.event_date, s.block_start_date) >= 3 THEN s.sequence_id + 1 ELSE s.sequence_id END AS sequence_id,
        -- 若序列ID更新，将当前日期设为新块的起始；否则沿用原块起始日期
        CASE WHEN DATEDIFF(g.event_date, s.block_start_date) >= 3 THEN g.event_date ELSE s.block_start_date END AS block_start_date
    FROM sequence_generator s
    JOIN grouped_sorted g 
        ON s.phone_number = g.phone_number 
        AND s.service = g.service 
        AND g.rn = s.rn + 1
)
-- 输出最终结果
SELECT phone_number, service, event_date, sequence_id
FROM sequence_generator
ORDER BY phone_number, service, event_date;

逻辑说明 #

1. 分组排序：先对每个phone_number+service分组的event_date排序，生成行号，确保递归能按时间顺序处理每一行。
2. 递归初始化：每个分组的第一行默认序列ID为1，同时将该行的日期作为第一个序列块的起始日期。
3. 滚动序列判断：对于后续每一行，计算当前日期与当前块起始日期的天数差：
   • 若差值≥3天：说明超出当前3天间隔，序列ID递增1，并将当前日期设为新块的起始日期
   • 若差值<3天：继续沿用当前序列ID和块起始日期

测试验证 #

用示例测试数据验证，输出完全符合预期：

测试数据 #

CREATE OR REPLACE TEMP VIEW test_data AS
SELECT '123' AS phone_number, 'A' AS service, DATE('2024-01-01') AS event_date
UNION ALL SELECT '123', 'A', '2024-01-02'
UNION ALL SELECT '123', 'A', '2024-01-04'
UNION ALL SELECT '123', 'A', '2024-01-05'
UNION ALL SELECT '123', 'B', '2024-01-01'
UNION ALL SELECT '123', 'B', '2024-01-03'
UNION ALL SELECT '123', 'B', '2024-01-06';

预期输出 #

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