如何编写函数批量聚合Pandas DataFrame？优化重复代码

场景背景 #

我们有一个记录国家销售数据的DataFrame，基于不同条件生成了多个子DataFrame，需要对每个子DataFrame的Sales和Profit字段执行聚合操作，避免重复编写代码。

原始数据与子DataFrame定义 #

import pandas as pd
import numpy as np

country_sales_list_of_lists = [
    ['Australia', 21421324, 342343, 'Pacific', 'Y'],
    ['England', 124233431, 43543464, 'Europe', 'Y'],
    ['Japan', 12431241341, 34267545, 'Asia', 'N'],
    ['India', 214343421, 342343, 'Asia', 'Y']
]

country_sales_df = pd.DataFrame(
    country_sales_list_of_lists,
    columns=['Country', 'Sales', 'Profit', 'Region', 'Otac_Group']
)

# 生成子DataFrame
otac_df = country_sales_df.query('Otac_Group == "Y"')
asia_df = country_sales_df.query('Region == "Asia"')
europe_df = country_sales_df.query('Region == "Europe"')
pacific_df = country_sales_df.query('Region == "Pacific"')

初始错误分析 #

最初的函数报错TypeError: 'DataFrameGroupBy' object is not callable，原因有两点：

1. 语法错误：groupby('Country')(['Sales','Profit'])写法错误，groupby返回的对象不能直接调用，正确写法是用方括号选取目标列：groupby('Country')[['Sales','Profit']]。
2. 逻辑错误：试图通过传入字符串df_name来命名新DataFrame，Python中无法直接将字符串作为变量名赋值，这种写法无效。

最简实现方案 #

方案1：字典批量处理（推荐） #

将所有需要聚合的子DataFrame存入字典，通过字典推导式批量生成聚合结果，既简洁又便于管理：

# 定义通用聚合函数
def aggregate_sales_profit(df):
    return df.groupby('Country')[['Sales', 'Profit']].agg([np.sum])

# 存储子DataFrame的字典（键为聚合结果的标识名）
target_dfs = {
    'pacific_agg': pacific_df,
    'europe_agg': europe_df,
    'asia_agg': asia_df,
    'otac_agg': otac_df
}

# 批量生成聚合结果
aggregated_results = {name: aggregate_sales_profit(df) for name, df in target_dfs.items()}

# 调用示例
print(aggregated_results['pacific_agg'])

方案2：直接基于原始DataFrame聚合（更高效） #

如果不需要保留单独的子DataFrame，可以直接在原始数据上按条件分组聚合，跳过子DataFrame的生成步骤：

# 按Region分组聚合所有区域数据
region_aggregates = country_sales_df.groupby(['Region', 'Country'])[['Sales', 'Profit']].agg(np.sum)
# 按Otac_Group分组聚合
otac_aggregates = country_sales_df.groupby(['Otac_Group', 'Country'])[['Sales', 'Profit']].agg(np.sum)

# 提取特定区域的聚合结果（比如Pacific）
pacific_agg = region_aggregates.xs('Pacific', level='Region')

方案3：简化原函数（保留子DataFrame变量） #

如果需要保留单独的聚合结果变量，可以简化原函数，去掉无用的df_name参数：

def aggregate_sales_profit(df):
    return df.groupby('Country')[['Sales', 'Profit']].agg([np.sum])

# 逐个生成聚合结果
pacific_df_agg = aggregate_sales_profit(pacific_df)
europe_df_agg = aggregate_sales_profit(europe_df)
asia_df_agg = aggregate_sales_profit(asia_df)
otac_df_agg = aggregate_sales_profit(otac_df)

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