Pandas大数据集按列值拆分多DataFrame的高效优化方案咨询

问题背景 #

需要将含437000行的Pandas DataFrame按第一列的28个唯一值拆分，每个唯一值对应一个DataFrame并存储在字典中。现有双重循环+逐行拼接的代码功能正常，但性能极差。

现有低效代码：

for i in range(size_unique):
    for j in range(size_data):
        if unique_values[i] == data_set.iloc[j, 0]:
            dict[unique_values[i]] = pd.concat([dict[unique_values[i]], data_set.iloc[j,:]])

性能问题根源 #

• 双重Python循环遍历所有唯一值和所有行，时间复杂度为O(M*N)（M为唯一值数量，N为总行数），Python循环本身效率远低于Pandas向量化操作。
• 每次调用pd.concat都会创建新的DataFrame对象，反复执行会产生大量不必要的内存开销和计算量。

优化方案 #

方案1：使用groupby（最优推荐）

Pandas原生的groupby方法是专门为分组场景设计的底层优化实现，性能远超手动循环：

# 按第一列（列索引0）分组，转换为字典
df_dict = dict(tuple(data_set.groupby(data_set.columns[0])))

# 若已知第一列列名（如示例中的'type'），写法更直观：
# df_dict = dict(tuple(data_set.groupby('type')))

执行后，df_dict的每个键对应第一列的唯一值，值为对应行组成的DataFrame。此方法全程为向量化操作，43万行数据可在毫秒级完成拆分。

方案2：布尔索引批量过滤（次优选择）

若不想用groupby，可通过布尔索引直接提取每个唯一值对应的所有行，避免逐行遍历：

unique_values = data_set.iloc[:, 0].unique()
df_dict = {}

for val in unique_values:
    # 布尔索引提取匹配行，copy避免后续修改原DataFrame
    df_dict[val] = data_set[data_set.iloc[:, 0] == val].copy()

此方法通过一次向量化过滤获取所有目标行，性能远优于原双重循环方案。

关于删除行的误区 #

匹配后删除行的思路不可取：删除行本身是O(N)操作，多次执行会导致性能下降，且索引变化易引发逻辑错误，完全没必要——上述两种方案已能高效完成需求，无需修改原DataFrame。

示例验证 #

针对提供的示例数据，使用groupby('type')后，df_dict['series_a']将返回如下DataFrame：

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