Python中高效拆分大型DataFrame为小DataFrame的最优方案

问题描述 #

我正在编写Python程序，需将数万行的大型DataFrame按列值拆分为多个小DataFrame，要求具备较高效率，以支持用户动态调整拆分方式并实时更新输出。

示例输入 #

示例输出 #

id=1的子DataFrame

id=2的子DataFrame

id=3的子DataFrame

朴素实现方式 #

我最初想到的实现方式如下，但认为这种方式会重复遍历数据，效率较低：

for id in list(df['id'].unique()):
    filtered_df = df[df['id'] == id]  

性能测试结果 #

我做了三种实现的性能测试，结果如下：

列表推导式过滤 #

%%timeit
[df.loc[df.id.eq(i)] for i in df.id.unique()]

9.96 ms ± 1.26 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

Groupby实现 #

%%timeit
dflist=[]
dflist2=[]
for  k,v  in df.groupby(['id']):
    var='id'+str(k)
    dflist.append(var)
    globals()[var] = v
    dflist2.append(v)

1.28 ms ± 92.2 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

字典推导式过滤 #

%%timeit
d = {id:df[df.id==id] for id in df.id.unique()}

9.19 ms ± 885 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

问题解答 #

1. 更高效的实现方法 #

使用groupby是最优选择，从测试结果也能看出它的效率远高于重复过滤的方式。你可以直接通过groupby获取分组后的子DataFrame，无需手动遍历唯一值并重复过滤：

# 获取分组后的字典，键是id值，值是对应的子DataFrame
grouped_dict = dict(df.groupby('id'))

# 或者直接遍历分组
for group_id, sub_df in df.groupby('id'):
    # 处理每个子DataFrame
    print(f"Group {group_id}:")
    print(sub_df)

2. 为什么groupby效率更高？ #

groupby的高效核心在于它只遍历原始DataFrame一次，而重复过滤的方式（比如循环每个唯一值做df[df.id == i]）会遍历DataFrameN次（N是唯一值的数量）。

具体来说：

• 重复过滤：每次过滤都要对整个DataFrame的id列进行比较操作，生成布尔掩码后再提取行，相当于重复扫描全量数据多次。
• Groupby：内部会先对id列进行排序或哈希映射，一次性将所有行分配到对应的分组中，后续直接返回分组结果，无需重复扫描数据。

此外，groupby的底层实现是基于Pandas的C扩展模块（如Cython），执行效率比纯Python循环+过滤要高得多，减少了Python层的开销。

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