Pandas按order_id分组筛选status为5的行执行过慢如何优化

原代码存在的问题 #

• 逻辑问题：你当前使用的groupby.apply+where逻辑不符合需求，只会保留status为5的行，其余行会被置为NaN，并不会保留符合条件order_id的所有行
• 性能问题：groupby.apply是Python层的循环操作，处理百万级数据时Python解释器的循环开销会被放大，导致耗时过长

优化方案 #

方案1：isin 筛选（最优，速度最快） #

这是性能最高的方案，全程使用矢量化操作，无额外分组开销，100万行数据处理耗时通常在10ms以内。
操作逻辑：先提取所有存在status=5的order_id集合，再直接筛选数据集里属于这个集合的所有行。
代码示例：

# 第一步：提取所有符合条件的order_id
valid_order_ids = order.loc[order['status'] == 5, 'order_id'].unique()
# 第二步：筛选所有对应order_id的行
order = order[order['order_id'].isin(valid_order_ids)]

方案2：groupby.transform 筛选 #

如果需要和分组逻辑强绑定，可以使用transform的矢量化实现，性能远高于apply，100万行数据处理耗时通常在50ms以内。
操作逻辑：按order_id分组后，标记所有存在status=5的分组，再筛选标记为真的行。
代码示例：

# 基础实现
order = order[order.groupby('order_id')['status'].transform(lambda x: (x == 5).any())]
# 无lambda的更优实现，性能更高
order = order[order.assign(flag=order['status'] == 5).groupby('order_id')['flag'].transform('max').astype(bool)]

以上两个方案对示例数据的输出结果一致，会保留id1、id2的全部8行数据，完全符合需求。

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