SQL LAG()函数的Pandas实现：计算用户会话间隔停留时长

原始DataFrame #

import pandas as pd

df = pd.DataFrame(
{'id': [10, 10, 10, 12, 12, 12, 12, 13, 13, 13],
 'session_id': [1, 3, 9, 1, 3, 5, 7, 1, 3, 5],
 'start_time': [5866, 6810, 8689, 8802, 8910, 9013, 9055, 9157, 9654, 9665],
 'end_time': [6808, 8653, 8722, 8881, 9001, 9049, 9062, 9651, 9659, 9725]
})

# 查看前5行
print(df.head())

输出：

id  session_id  start_time  end_time
0   10           1        5866      6808
1   10           3        6810      8653
2   10           9        8689      8722
3   12           1        8802      8881
4   12           3        8910      9001

需求说明 #

需要新增stay_time列，用于存储用户当前会话结束后到下一会话开始前的停留时长，期望输出如下：

id  session_id  start_time  end_time  stay_time
0   10           1        5866      6808          0
1   10           3        6810      8653          2
2   10           9        8689      8722         36
3   12           1        8802      8881          0
4   12           3        8910      9001         29
5   12           5        9013      9049         12
6   12           7        9055      9062          6
7   13           1        9157      9651          0
8   13           3        9654      9659          3
9   13           5        9665      9725          6

SQL等价实现 #

假设数据表名为participants，对应的SQL语句如下：

select p.*,
    start_time - lag(end_time, 1, start_time) over(partition by id order by session_id) as stay_time
from participants p

Pandas实现方案 #

方法一：分组+移位+填充（推荐） #

# 按用户id分组，获取每组内上一行的end_time，无上一行时用当前start_time填充
prev_end = df.groupby('id')['end_time'].shift(1).fillna(df['start_time'])

# 计算停留时长：当前会话开始时间 - 上一会话结束时间
df['stay_time'] = (df['start_time'] - prev_end).astype(int)

print(df)

方法二：分组应用lambda函数 #

# 分组后对每组单独计算停留时长
df['stay_time'] = df.groupby('id').apply(
    lambda x: x['start_time'] - x['end_time'].shift(1).fillna(x['start_time'])
).reset_index(level=0, drop=True).astype(int)

print(df)

逻辑对应说明 #

• groupby('id')：对应SQL的partition by id，按用户维度分组处理
• shift(1)：对应SQL的lag(end_time, 1)，取分组内上一条记录的结束时间
• fillna(df['start_time'])：对应SQL的lag函数默认值参数，用户首个会话无前置记录时用自身开始时间填充，保证计算结果为0
• 最终通过start_time - prev_end得到两次会话间的停留时长

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