如何在SQL/Python中按user_id生成重置式attempt列？

需求说明 #

现有数据集包含user_id、status、date字段，需新增attempt列，规则为：

• 按user_id独立分组计算
• 每当用户出现completed状态后，后续的尝试计数重置为1
• 同一尝试周期内（从上次completed后第一条记录到下一条completed），attempt按日期顺序递增
• 数据集规模较大，需低计算量的实现方案

SQL实现方案（主流数据库通用） #

利用窗口函数实现矢量化计算，避免关联查询，计算复杂度为O(n)，适合大数据量：

WITH reset_markers AS (
    SELECT 
        user_id,
        status,
        date,
        -- 标记新尝试周期的起点：首次记录 或 前一条记录为completed
        CASE 
            WHEN LAG(status) OVER (PARTITION BY user_id ORDER BY date) = 'completed' THEN 1
            WHEN ROW_NUMBER() OVER (PARTITION BY user_id ORDER BY date) = 1 THEN 1
            ELSE 0
        END AS is_new_attempt
    FROM your_table
),
attempt_groups AS (
    SELECT 
        user_id,
        status,
        date,
        -- 累积标记生成尝试周期ID
        SUM(is_new_attempt) OVER (PARTITION BY user_id ORDER BY date) AS group_id
    FROM reset_markers
)
SELECT 
    user_id,
    status,
    date,
    -- 每个周期内的行号即为attempt值
    ROW_NUMBER() OVER (PARTITION BY user_id, group_id ORDER BY date) AS attempt
FROM attempt_groups
ORDER BY user_id, date;

Python实现方案（Pandas矢量化操作） #

采用Pandas的分组矢量化计算，避免逐行循环，适合大规模数据集：

import pandas as pd

# 加载数据并预处理日期格式
df = pd.read_csv("your_data.csv")
df["date"] = pd.to_datetime(df["date"])

# 按用户和日期排序，确保计算顺序正确
df = df.sort_values(["user_id", "date"])

# 标记每个新尝试周期的起点
df["is_new_attempt"] = df.groupby("user_id")["status"].shift(1).eq("completed").fillna(True).astype(int)

# 生成每个用户的尝试周期分组ID
df["group_id"] = df.groupby("user_id")["is_new_attempt"].cumsum()

# 计算每个周期内的attempt值
df["attempt"] = df.groupby(["user_id", "group_id"]).cumcount() + 1

# 清理中间辅助列
df = df.drop(["is_new_attempt", "group_id"], axis=1)

# 输出结果
print(df)

问题排查说明 #

之前多用户场景失效的核心原因是未按user_id做独立分组计算，上述方案均明确通过PARTITION BY user_id（SQL）或groupby('user_id')（Python）确保每个用户的计数逻辑独立隔离。

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