如何优化pandas按条件合并行填充缺失起止时间的代码

业务需求 #

待处理数据集包含start_time、end_time、id、url四个核心字段，同一个id和url分组下的时间记录分散在不同行，需要按以下规则补全缺失的时间值：

• 若某行存在有效start_time但end_time为空，用大于等于该start_time的最近end_time填充空值，同时删除被匹配到的end_time所在原行
• 所有带有效start_time的行处理完成后，若仍有start_time为空的行，直接将该行end_time赋值给start_time
• 若某start_time未匹配到符合条件的end_time，直接将该行start_time赋值给end_time

现有实现代码 #

def process(self, param, context):
    df = context['data']
    # df = df.drop_duplicates()
    key_cols = param['keys_cols']
    start_time_col = param['start_time_col']
    end_time_col = param['end_time_col']
    guid_col = param.get('guid_col','guid')
    df_groupby = df.groupby(key_cols).size().reset_index()
    final_dfs = []
    condition = ''
    for key in key_cols:
        if condition == '':
            condition = '(df[\''+str(key)+'\']==row[\''+str(key)+'\'])'
        else:
            condition = condition + ' & ' +'(df[\'' + str(key) + '\']==row[\'' + str(key) + '\'])'
    for index, row in df_groupby.iterrows():
        sub_df = df[eval(condition)]
        if sub_df[start_time_col].isnull().sum() != len(sub_df[start_time_col]) and (sub_df[end_time_col].isnull().sum() != len(sub_df[end_time_col])):
            sub_df = sub_df.sort_values([start_time_col, end_time_col], ascending=True)
            subdf_start_time_not_null = sub_df[sub_df[start_time_col].notnull()]
            subdf_end_time_not_null = sub_df[sub_df[end_time_col].notnull()]
            subdf_end_time_not_null['combined'] = subdf_end_time_not_null[end_time_col] +"__"+ subdf_end_time_not_null[guid_col]
            end_time_values = subdf_end_time_not_null['combined'].values.tolist()
            for row_number, (stime_index, stime_row) in enumerate(subdf_start_time_not_null.iterrows()):
                delete_index = row_number
                if row_number < len(end_time_values):
                    end_time_value = np.nan
                    if int(str(subdf_start_time_not_null.at[stime_index,start_time_col]).replace(":","").replace(" ","").replace("-","")) <= int(str(end_time_values[row_number]).split("__")[0].replace(":","").replace(" ","").replace("-","")):
                        end_time_value = end_time_values[row_number]
                        subdf_start_time_not_null.at[stime_index,end_time_col] = str(end_time_values[row_number]).split("__")[0]
                    else:
                        prev_index = end_time_values.index(end_time_values[row_number])
                        for end_time in end_time_values:
                            current_index = end_time_values.index(end_time)
                            if current_index > prev_index:
                                if int(str(subdf_start_time_not_null.at[stime_index,start_time_col]).replace(":","").replace(" ","").replace("-","")) <= int(str(end_time_values[current_index]).split("__")[0].replace(":","").replace(" ","").replace("-","")):
                                    subdf_start_time_not_null.at[stime_index, end_time_col] = end_time_values[current_index]
                                    delete_index = current_index
                                    end_time_value = end_time_values.pop(delete_index)
                                    break
                    subdf_end_time_not_null = subdf_end_time_not_null[subdf_end_time_not_null[guid_col]!=end_time_value.split("__")[1]]
                else:
                    subdf_start_time_not_null.at[stime_index,end_time_col] = subdf_start_time_not_null.at[stime_index,start_time_col]
            subdf_end_time_not_null.drop('combined', axis=1, inplace=True)
            sub_df = pd.concat([subdf_start_time_not_null,subdf_end_time_not_null])
        sub_df[start_time_col] = np.where(sub_df[start_time_col].isnull(),sub_df[end_time_col],sub_df[start_time_col])
        sub_df[end_time_col] = np.where(sub_df[end_time_col].isnull(),sub_df[start_time_col],sub_df[end_time_col])
        final_dfs.append(sub_df)
    df = pd.concat(final_dfs)
    context['data'] = df
    context['continue'] = True
    return context

调用参数示例 #

param = {"keys_cols":['id', 'url'], "start_time_col":"start_time","end_time_col":"end_time"}

优化建议 #

原有代码性能差的核心原因是大量使用循环、动态eval执行、重复字符串转时间的冗余计算，可从以下几个方向优化：

1. 替换动态分组过滤逻辑：原有代码先分组求size再循环每一组用eval过滤子df，开销极高且不安全，直接使用pandas原生的groupby(key_cols).apply()处理每个分组，省去eval和逐行过滤的开销。
2. 提前转换时间格式：原有代码每次比较时间都要把字符串转成数字，重复计算量极大，在处理前统一把start_time和end_time转成pandas datetime类型，直接比较大小，性能提升可达数十倍。
3. 用二分查找替代循环匹配：每一组内将非空的end_time排序后，用bisect模块的bisect_left方法找第一个大于等于当前start_time的end_time位置，匹配复杂度从O(n)降到O(logn)，匹配到后直接移除对应位置的end_time即可，不需要额外维护guid过滤逻辑。
4. 减少临时对象拷贝：原有代码反复拼接、过滤子df产生大量临时内存拷贝，直接在分组内的子df上做修改，减少不必要的内存操作。

优化后核心逻辑示例 #

import pandas as pd
import bisect

def process_group(sub_df, start_col, end_col):
    # 拆分起始时间行和结束时间行
    start_rows = sub_df[sub_df[start_col].notna()].sort_values(start_col).reset_index(drop=True)
    end_rows = sub_df[sub_df[end_col].notna()].sort_values(end_col).reset_index(drop=True)
    
    end_times = end_rows[end_col].tolist()
    used_end_idx = set()
    
    # 匹配最近的end time
    for i, st in enumerate(start_rows[start_col]):
        idx = bisect.bisect_left(end_times, st)
        while idx < len(end_times) and idx in used_end_idx:
            idx += 1
        if idx < len(end_times):
            start_rows.at[i, end_col] = end_times[idx]
            used_end_idx.add(idx)
        else:
            start_rows.at[i, end_col] = st
    
    # 未被匹配的end时间行
    unused_end_rows = end_rows[~end_rows.index.isin(used_end_idx)].copy()
    # 补全start为空的行
    unused_end_rows[start_col] = unused_end_rows[end_col]
    
    return pd.concat([start_rows, unused_end_rows], ignore_index=True)

def process(self, param, context):
    df = context['data']
    key_cols = param['keys_cols']
    start_col = param['start_time_col']
    end_col = param['end_time_col']
    
    # 提前转时间格式
    df[start_col] = pd.to_datetime(df[start_col], errors='coerce')
    df[end_col] = pd.to_datetime(df[end_col], errors='coerce')
    
    # 分组处理
    df = df.groupby(key_cols, group_keys=False).apply(
        lambda x: process_group(x, start_col, end_col)
    ).reset_index(drop=True)
    
    # 兜底补全所有空值
    df[start_col] = df[start_col].fillna(df[end_col])
    df[end_col] = df[end_col].fillna(df[start_col])
    
    context['data'] = df
    context['continue'] = True
    return context

优化后的代码处理7万行数据耗时可从分钟级降到秒级。

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