PySpark SQL实现考试尝试编号（含6个月重置逻辑）

核心逻辑是先通过时间间隔判断分组边界，再在每个分组内生成递增的尝试编号，具体实现步骤如下：

1. 基础排序与前置时间获取 #

假设你的考试记录表名为exam_records，包含字段：user_id（用户ID）、exam_id（考试ID）、completion_time（考试完成时间，timestamp类型）。首先按用户和考试分组，按完成时间排序，同时获取每条记录的上一条完成时间。

2. 标记分组重置边界 #

计算当前记录与上一条记录的月份间隔，若间隔超过6个月，或者是该用户该考试的第一条记录，则标记为新分组的起点。通过累加这些标记，生成每个记录所属的分组ID。

3. 生成递增尝试编号 #

在每个用户-考试-分组内，用行号函数生成从1开始递增的AttemptNumber。

完整PySpark SQL代码 #

WITH ranked_records AS (
    SELECT
        user_id,
        exam_id,
        completion_time,
        -- 获取同一用户同一考试的上一条完成时间
        LAG(completion_time) OVER (PARTITION BY user_id, exam_id ORDER BY completion_time) AS prev_completion_time,
        -- 标记是否为新分组起点：第一条记录 或 与上一条间隔超6个月
        CASE
            WHEN LAG(completion_time) OVER (PARTITION BY user_id, exam_id ORDER BY completion_time) IS NULL THEN 1
            WHEN MONTHS_BETWEEN(completion_time, LAG(completion_time) OVER (PARTITION BY user_id, exam_id ORDER BY completion_time)) > 6 THEN 1
            ELSE 0
        END AS is_new_group
    FROM exam_records
),
grouped_records AS (
    SELECT
        *,
        -- 累加分组标记，生成唯一分组ID
        SUM(is_new_group) OVER (PARTITION BY user_id, exam_id ORDER BY completion_time ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS attempt_group_id
    FROM ranked_records
)
SELECT
    user_id,
    exam_id,
    completion_time,
    -- 在每个分组内生成尝试编号
    ROW_NUMBER() OVER (PARTITION BY user_id, exam_id, attempt_group_id ORDER BY completion_time) AS AttemptNumber
FROM grouped_records
ORDER BY user_id, exam_id, completion_time;

代码关键部分说明 #

• LAG(completion_time)：拉取同一用户同一考试的上一条完成记录时间，用于计算时间间隔。
• MONTHS_BETWEEN：精确计算两个时间的月份差，作为判断是否触发编号重置的依据。
• SUM(is_new_group) OVER (...)：通过累加新分组标记，将连续间隔不超过6个月的记录归为同一组，间隔超6个月的则开启新分组。
• ROW_NUMBER()：在每个分组内按完成时间排序，生成从1开始的递增尝试编号。

示例验证 #

假设某用户某考试的记录如下：

运行查询后结果：

可以看到，2023-02-10的记录与上一条间隔超过6个月，尝试编号重置为1，符合需求。

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