优化pandas DataFrameGroupBy使用：按日期产品分组计算累计容量

问题背景 #

现有大型DataFrame结构如下：

需求：按DATE和PRODUCT分组，每组按PRICE升序排序后，计算CAPACITY的累计值，期望结果如下：

当前实现方式（效率较低）：

df_result = pd.DataFrame()
for i, group in df.groupby(by=['DATE', 'PRODUCT']):
    group.sort_values('PRICE', inplace=True)
    group['CUMULATIVE'] = group['CAPACITY'].cumsum()
    df_result = pd.concat([df_result, group], ignore_index=True)

高效优化方案 #

核心思路 #

利用Pandas的向量化操作替代Python循环，避免多次concat带来的内存开销，同时确保分组内排序后再计算累计值。

步骤1：处理PRICE列的数值类型 #

由于原始PRICE是带逗号的字符串，需先转换为数值类型，否则字符串排序会出现逻辑错误：

df['PRICE'] = df['PRICE'].str.replace(',', '.').astype(float)

步骤2：全局排序+分组累计求和 #

先按分组键（DATE、PRODUCT）和PRICE升序排序，再分组计算CAPACITY的累计值：

# 按分组键+PRICE升序排序，确保同组内数据按价格排序
df_sorted = df.sort_values(by=['DATE', 'PRODUCT', 'PRICE'])

# 分组计算累计值，Pandas会自动在每组内独立计算cumsum
df_sorted['CUMULATIVE'] = df_sorted.groupby(['DATE', 'PRODUCT'])['CAPACITY'].cumsum()

# 可选：重置索引
df_result = df_sorted.reset_index(drop=True)

方案优势 #

• 性能提升：向量化操作由Pandas底层优化实现，比Python循环快数倍，尤其适合大型数据集；
• 代码简洁：避免循环和多次concat，代码可读性更高；
• 稳定性强：减少手动处理分组时可能出现的索引错误或数据拼接问题。

备选方案（apply方式） #

如果需要更灵活的组内处理，可使用groupby.apply，但性能略逊于上述方案：

df['PRICE'] = df['PRICE'].str.replace(',', '.').astype(float)

df_result = (
    df.groupby(['DATE', 'PRODUCT'], group_keys=False)
    .apply(lambda x: x.sort_values('PRICE').assign(CUMULATIVE=x['CAPACITY'].cumsum()))
    .reset_index(drop=True)
)

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