如何用Pythonic的Pandas方法高效按用户划分考试时段并优化性能？

规则说明 #

• 首题归为时段1
• 满足以下任一条件时，时段序号+1：
  1. 当前题开始时间与上一题结束时间间隔超过10分钟
  2. 单题耗时超过15分钟
• 不满足上述条件时，时段序号保持不变

现有实现及性能 #

1. 循环遍历方案：处理2000用户耗时11秒

def user_exam_periods(user,gap_cutoff,length_cutoff,data_df=data_df):
    # 获取当前用户数据
    user_df = data_df[data_df['rater_email']==user]
    # 按开始时间排序
    user_df = user_df.sort_values('start_time')
    # 添加空的时段列
    user_df['exam_period'] = ''
    # 重置索引以便循环中使用.loc
    user_df = user_df.reset_index()
    user_df.loc[0,'exam_period'] = 1
    # 时段递增逻辑
    for i in range(1,user_df.shape[0]):
        if user_df.loc[i,'start_time'] > user_df.loc[i-1,'end_time'] + pd.Timedelta(minutes=gap_cutoff):
            user_df.loc[i,'exam_period'] = user_df.loc[i-1,'exam_period'] + 1
        elif user_df.loc[i,'start_time'] < user_df.loc[i,'end_time'] - pd.Timedelta(minutes=length_cutoff):
            user_df.loc[i,'exam_period'] = user_df.loc[i-1,'exam_period'] + 1
        else:
            user_df.loc[i,'exam_period'] = user_df.loc[i-1,'exam_period']
    
    # 恢复原索引用于后续合并
    user_df.set_index('index')
    return user_df[['rater_email','task_id','exam_period']]

2. Groupby向量化方案：处理2000用户耗时8.5秒，但仍无法满足未来大规模用户的性能需求

def calculate_exam_periods(data_df, gap_cutoff, length_cutoff):
    data_df = data_df.sort_values(['rater_email', 'start_time'])

    # 计算行间时间差
    time_diff = data_df.groupby('rater_email')['start_time'].diff()

    # 标记间隔过长或单题耗时过长的情况
    gaps = time_diff > pd.Timedelta(minutes=gap_cutoff)
    long_tasks = data_df['end_time'] - data_df['start_time'] > pd.Timedelta(minutes=length_cutoff)

    # 累计触发条件并生成时段序号
    data_df['exam_period'] = (gaps | long_tasks).astype(int).groupby(data_df['rater_email']).cumsum() + 1

    return data_df[['rater_email', 'task_id', 'exam_period']]

优化方案：全向量化+预转换时间类型 #

以下方案通过减少重复计算、预转换时间类型、避免不必要的中间变量进一步提升性能，适合大规模数据场景：

def optimize_exam_periods(data_df, gap_cutoff=10, length_cutoff=15):
    # 确保时间列是datetime类型（提前转换一次，避免重复计算）
    data_df = data_df.copy()
    data_df['start_time'] = pd.to_datetime(data_df['start_time'])
    data_df['end_time'] = pd.to_datetime(data_df['end_time'])
    
    # 按用户和开始时间排序（仅排序一次）
    data_df = data_df.sort_values(['rater_email', 'start_time'], ignore_index=True)
    
    # 预计算时间阈值（转换为Timedelta一次）
    gap_threshold = pd.Timedelta(minutes=gap_cutoff)
    length_threshold = pd.Timedelta(minutes=length_cutoff)
    
    # 计算上一题的结束时间（按用户分组后shift）
    prev_end = data_df.groupby('rater_email')['end_time'].shift(1)
    
    # 标记需要递增时段的条件
    need_increment = (
        (data_df['start_time'] - prev_end > gap_threshold) |
        (data_df['end_time'] - data_df['start_time'] > length_threshold)
    )
    # 处理每组的第一个元素（初始为False，不影响cumsum）
    need_increment = need_increment.fillna(False)
    
    # 按用户累计递增条件，生成时段序号
    data_df['exam_period'] = need_increment.groupby(data_df['rater_email']).cumsum() + 1
    
    return data_df[['rater_email', 'task_id', 'exam_period']]

优化点说明

• 提前统一转换时间列为datetime类型，避免后续重复转换开销
• 预计算时间阈值，减少pd.Timedelta的重复实例化
• 使用shift()替代diff()计算上一题结束时间，逻辑更直观且减少中间变量
• 合并条件判断，减少布尔数组的创建次数
• 使用ignore_index=True避免后续索引操作的额外开销

示例数据 #

rater_email,locale,task_id,start_time,end_time
foo@bar.com,en_US,foobar1,2023-03-13T19:31:05Z,2023-03-13T19:42:15Z
foo@bar.com,en_US,foobar2,2023-03-13T19:42:17Z,2023-03-13T19:47:12Z
foo@bar.com,en_US,foobar3,2023-03-13T19:47:15Z,2023-03-13T20:01:46Z
foo@bar.com,en_US,foobar4,2023-03-13T20:01:48Z,2023-03-13T20:07:37Z
foo@bar.com,en_US,foobar5,2023-03-13T20:07:39Z,2023-03-13T20:20:18Z
foo@bar.com,en_US,foobar6,2023-03-13T20:20:21Z,2023-03-13T20:29:44Z
foo@bar.com,en_US,foobar7,2023-03-13T20:31:01Z,2023-03-13T20:35:32Z
foo@bar.com,en_US,foobar8,2023-03-13T20:35:35Z,2023-03-13T20:41:35Z
foo@bar.com,en_US,foobar9,2023-03-13T20:41:37Z,2023-03-13T20:46:03Z
foo@bar.com,en_US,foobar10,2023-03-13T20:46:05Z,2023-03-13T20:51:32Z
foo@bar.com,en_US,foobar11,2023-03-13T20:53:21Z,2023-03-13T21:01:21Z
foo@bar.com,en_US,foobar12,2023-03-13T21:01:23Z,2023-03-13T21:17:52Z
foo@bar.com,en_US,foobar13,2023-03-13T21:17:54Z,2023-03-13T21:24:18Z
foo@bar.com,en_US,foobar14,2023-03-13T21:25:41Z,2023-03-13T21:31:24Z
foo@bar.com,en_US,foobar15,2023-03-13T21:31:26Z,2023-03-13T21:35:14Z
foo@bar.com,en_US,foobar16,2023-03-13T21:35:17Z,2023-03-13T21:42:14Z
foo@bar.com,en_US,foobar17,2023-03-13T21:42:16Z,2023-03-13T21:48:33Z
foo@bar.com,en_US,foobar18,2023-03-13T21:48:35Z,2023-03-13T21:50:53Z
foo@bar.com,en_US,foobar19,2023-03-13T21:50:55Z,2023-03-13T21:58:18Z
foo@bar.com,en_US,foobar20,2023-03-13T21:58:21Z,2023-03-13T22:05:38Z
foo@bar.com,en_US,foobar21,2023-03-13T22:05:40Z,2023-03-13T22:18:00Z
foo@bar.com,en_US,foobar22,2023-03-13T22:18:02Z,2023-03-13T22:25:27Z
foo@bar.com,en_US,foobar23,2023-03-13T22:25:29Z,2023-03-13T22:32:20Z
foo@bar.com,en_US,foobar24,2023-03-13T22:32:23Z,2023-03-13T22:42:27Z
foo@bar.com,en_US,foobar25,2023-03-13T22:42:29Z,2023-03-13T22:45:47Z
foo@bar.com,en_US,foobar26,2023-03-13T22:45:49Z,2023-03-13T22:49:25Z
foo@bar.com,en_US,foobar27,2023-03-13T22:49:27Z,2023-03-13T22:53:06Z
foo@bar.com,en_US,foobar28,2023-03-13T22:53:08Z,2023-03-13T23:10:07Z
foo@bar.com,en_US,foobar29,2023-03-13T23:11:27Z,2023-03-13T23:17:08Z
foo@bar.com,en_US,foobar30,2023-03-13T23:17:10Z,2023-03-13T23:23:25Z

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