Python函数中DataFrame分组聚合后用pd.concat拼接如何修改原入参

原有代码问题梳理 #

• 分组键不匹配：需求是对B2C数据按col4、col5分组，原代码错误使用col1、col2作为分组键，导致聚合结果不符合预期
• 入参修改逻辑错误：Python中函数参数是赋值传递，直接对局部变量output_data重新赋值只会修改局部变量的指向，不会改变外部传入的原DataFrame对象内容

修正后实现代码 #

import pandas as pd

def fun(output_data):
    # 构造聚合规则，可根据实际需求调整需要sum的列
    dic_ = dict.fromkeys(output_data.columns, "first")
    # 分组键不需要配置聚合规则，所以移除对应配置
    dic_.pop("col4")
    dic_.pop("col5")
    # 示例中col3需要求和，所以加入update字典，不需要可删除
    dic_.update({
        'col3': "sum",
        'col9': "sum",
        'col10': "sum",
        'col11': "sum",
        'col12': "sum",
    })

    # 筛选B2C数据按col4、col5分组聚合
    b2c_data = output_data[output_data['col100'].eq('B2C')]
    tmp = b2c_data.groupby(
        ['col4', 'col5'], sort=False, as_index=False
    ).agg(dic_)[list(output_data.columns)].reset_index(drop=True)

    # 拼接聚合后的B2C数据和非B2C数据
    non_b2c_data = output_data[output_data['col100'].ne('B2C')]
    res = pd.concat([tmp, non_b2c_data], ignore_index=True)

    # 直接修改原入参DataFrame的内容，不重新绑定变量
    output_data.drop(output_data.index, inplace=True)
    output_data[res.columns] = res

验证说明 #

用你给出的示例输入调用函数后，外部传入的原DataFrame会直接被修改为预期输出结果，不需要额外接收返回值。

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