Pandas分组高效筛选：50万+数据按自定义条件过滤组内记录

需求说明 #

现有包含50万+条记录的Pandas DataFrame，需按first_roll_up、date、granular_timestamp列分组，执行以下筛选逻辑：

• 若组内存在top列有值的记录，仅保留该记录；
• 若组内无top列有值的记录，则保留组内所有记录。

输入示例 #

输出示例 #

问题背景 #

此前尝试使用以下代码处理，但自定义函数执行耗时超10分钟，改用transform优化也无效果：

df.groupby(['first_roll_up', 'sub', 'top', 'date', 'granular_timestamp'], sort=False)
        .apply(custom_function_to_filter_each_group_records)

高效解决方案 #

核心思路 #

避免使用apply逐组处理（低效），改用向量化分组标记+布尔索引筛选，全程利用Pandas的内置优化，处理百万级数据仅需数秒。

代码实现 #

import pandas as pd

# 1. 标记每条记录所在组是否存在top列非空的记录
df['has_top_in_group'] = df.groupby(['first_roll_up', 'date', 'granular_timestamp'])['top'].transform(
    lambda x: x.notna().any()
)

# 2. 筛选：要么当前记录top非空，要么所在组无top非空记录
df_filtered = df[(df['top'].notna()) | (~df['has_top_in_group'])]

# 可选：删除临时标记列
df_filtered = df_filtered.drop(columns=['has_top_in_group'])

方案说明 #

1. 分组标记：通过transform为每条记录生成所在组的has_top_in_group标记，该操作是向量化的，远快于逐组apply。
2. 布尔筛选：直接用布尔索引过滤数据，这是Pandas中效率最高的筛选方式之一。
3. 优化点：原代码的groupby列错误（包含了sub和top，导致分组粒度完全不符合需求），正确的分组列应为需求指定的first_roll_up、date、granular_timestamp。

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