优化百万级Pandas DataFrame：保留sub_batch首个done状态并提速

问题核心 #

你有一个500万行的Pandas DataFrame，需实现：每个sub_batch首次出现"done"状态后，删除该sub_batch后续所有"done"行，保留其余行并维持原排序；未出现"done"的sub_batch数据完全保留。原代码因双重循环遍历，耗时长达2天，需要高效优化方案。

原代码低效原因 #

原代码采用双重循环遍历所有sub_batch和DataFrame行，时间复杂度为O(N*M)（N为总行数，M为唯一sub_batch数量），面对百万级数据时性能极差，必须用Pandas的矢量化分组操作替代循环。

优化方案（矢量化实现） #

以下两种方案均基于Pandas内置的分组运算，处理500万行数据仅需数分钟：

方案一：用累积计数快速筛选

import pandas as pd

# 1. 标记所有"done"行
df['is_done'] = df['output'] == 'done'

# 2. 对每个sub_batch，计算done行的累积出现次数
df['done_cumulative'] = df.groupby('sub_batch')['is_done'].cumsum()

# 3. 生成筛选掩码：保留非done行，或第一个出现的done行
keep_mask = (~df['is_done']) | ((df['done_cumulative'] == 1) & df['is_done'])

# 4. 筛选数据并清理临时列
df_filtered = df[keep_mask].drop(columns=['is_done', 'done_cumulative'])

原理说明：

• done_cumulative对每个sub_batch内的"done"行累加计数，第一个"done"行的计数值为1，后续"done"行计数值≥2
• 筛选逻辑直接通过掩码实现：非"done"行全部保留；仅保留计数值为1的"done"行（即每个sub_batch的第一个"done"）

方案二：分组定位首次done位置（更直观）

import pandas as pd

# 1. 标记所有"done"行
df['is_done'] = df['output'] == 'done'

# 2. 定义分组过滤函数：保留非done行或第一个done行
def filter_group(group):
    first_done_idx = group[group['is_done']].index.min()
    if pd.isna(first_done_idx):
        # 该sub_batch无done行，全保留
        return pd.Series(True, index=group.index)
    # 保留条件：非done行 或 是第一个done行
    return (~group['is_done']) | (group.index == first_done_idx)

# 3. 应用分组过滤，生成掩码
keep_mask = df.groupby('sub_batch').apply(filter_group).reset_index(level=0, drop=True)

# 4. 筛选并清理
df_filtered = df[keep_mask].drop(columns=['is_done'])

原理说明：

• 对每个sub_batch分组，先找到第一个"done"行的索引
• 若分组内无"done"，则全保留；否则保留非"done"行和第一个"done"行

性能对比 #

原双重循环方案：时间与sub_batch数量正相关，500万行需数十小时
优化后的矢量化方案：时间与总行数线性相关，500万行仅需数分钟即可完成

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