pandas中groupby().agg()无法保留DataFrame子类类型的问题咨询

这不是bug，是pandas中agg方法的内部实现逻辑导致的：普通聚合方法（如sum）会继承groupby对象关联的子类构造器，但agg在处理字典式参数这类复杂聚合场景时，内部结果构造路径没有传递自定义子类的构造器，最终返回原生DataFrame。

解决方法 #

方法1：重写自定义DataFrame的groupby方法，返回自定义GroupBy子类并修正agg逻辑 #

通过自定义DataFrameGroupBy子类，重写agg方法，将结果显式转换为自定义MyDataFrame类型：

import pandas as pd
from pandas.core.groupby.generic import DataFrameGroupBy

class MyDataFrameGroupBy(DataFrameGroupBy):
    def agg(self, *args, **kwargs):
        # 调用原生agg逻辑获取结果
        result = super().agg(*args, **kwargs)
        # 转换为自定义DataFrame子类
        return MyDataFrame(result)

class MySeries(pd.Series):
    pass

class MyDataFrame(pd.DataFrame):
    @property
    def _constructor(self):
        return MyDataFrame
    _constructor_sliced = MySeries
    
    def groupby(self, *args, **kwargs):
        # 调用原生groupby获取分组对象
        gb = super().groupby(*args, **kwargs)
        # 包装为自定义GroupBy子类
        return MyDataFrameGroupBy(
            gb.obj, gb.grouper, gb.axis, gb.level, gb.sort, gb.group_keys
        )

MySeries._constructor_expanddim = MyDataFrame

# 测试
df = MyDataFrame({"a": reversed(range(10)), "b": list('aaaabbbccc')})
print(type(df.groupby("b").sum()))
# <class '__main__.MyDataFrame'>
print(type(df.groupby("b").agg({"a": "sum"})))
# <class '__main__.MyDataFrame'>

方法2：手动转换agg结果（简易版） #

如果不需要全局适配，也可以在agg调用后直接用自定义子类包裹结果：

result = MyDataFrame(df.groupby("b").agg({"a": "sum"}))
print(type(result))
# <class '__main__.MyDataFrame'>

补充说明 #

pandas在处理agg的复杂参数时，由于返回结构可能更灵活（比如多列多聚合函数），内部默认使用原生DataFrame构造器来保证兼容性，这是设计上的取舍而非bug。通过自定义GroupBy子类的方式，可以在不修改pandas源码的前提下，让agg结果保留自定义子类类型。

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