如何减少200万行数据集按条件查行时嵌套循环的计算耗时

原代码性能瓶颈分析 #

• 嵌套for循环遍历仓库和SKU，完全没有利用pandas的向量化运算能力，属于低效率的行级迭代
• 每次调用append方法都会生成全新的DataFrame，随着结果集变大，内存拷贝开销会指数级增长
• 重复的布尔索引筛选操作产生了大量冗余计算，进一步拉长了执行时间

优化后代码 #

方案1：仅取每组第一条最快记录（适合无并列耗时的场景，性能最优） #

通过groupby + idxmin直接定位每个仓库、每个产品分组下耗时最小的行索引，一步提取结果，200万行数据通常可以在几秒到几十秒内执行完成：

def model_selection_df(df):
    # 按【仓库编码、产品编码】分组，取每组耗时最小的行索引
    min_time_idx = df.groupby(['WAREHOUSE_CODE', 'VEND_PART'])['TIME_TAKEN'].idxmin()
    # 直接提取对应行，重置索引后返回
    return df.loc[min_time_idx].reset_index(drop=True)

方案2：保留所有并列最快的运输方式（适合存在多个运输方式耗时同为最小值的场景） #

如果需要把同一个分组下所有耗时等于最小值的运输方式都保留，可以用groupby + transform生成筛选掩码：

def model_selection_df(df):
    # 生成布尔掩码：标记每行是否属于所属分组的耗时最小值
    is_min_time = df['TIME_TAKEN'] == df.groupby(['WAREHOUSE_CODE', 'VEND_PART'])['TIME_TAKEN'].transform('min')
    # 筛选符合条件的行，重置索引后返回
    return df[is_min_time].reset_index(drop=True)

额外优化建议 #

• 可以提前把WAREHOUSE_CODE、VEND_PART字段转为category类型，减少内存占用的同时进一步提升groupby的执行速度
• 后续避免在pandas中使用for循环遍历分组做计算，内置的groupby算子是C语言实现的，性能比手写Python循环高2~3个数量级
• 不要在循环中反复调用DataFrame的append方法，这类操作的时间复杂度是O(n²)，数据量越大效率越低

内容的提问来源于stack exchange，提问作者Rajat kumar Pal