继承pd.DataFrame的类自pickle失效，求原因及解决方案

问题背景 #

想要在继承pd.DataFrame的类中添加name属性存储元数据，原生DataFrame.name无法在pickle序列化后保留，且暂不考虑实验性的df.attr['name']方案。重写to_pickle和read_pickle方法后，反序列化得到的对象丢失了name属性。

初始实现代码 #

import pandas as pd
import pickle

class NamedDataFrame(pd.DataFrame):
    '''
    a dataframe with a name
    '''
    def __init__(self, data=None, index=None, columns=None, dtype=None, copy=False, name: str = None):
        super().__init__(data, index, columns, dtype, copy)
        self.name = name
    
    # 尝试重写pickle方法以保存name属性
    def to_pickle(self, path, compression='infer', protocol=4):
        print("pickling myself")
        with open(path, 'wb') as f:
            pickle.dump(self, f, protocol)

    @classmethod
    def read_pickle(cls, path):
        with open(path, 'rb') as f:
            return pickle.load(f)

问题复现 #

>>> ndf = NamedDataFrame(data=mydf,name='mytestname')
>>> ndf.name
mytestname

>>> ndf.to_pickle(mypath)
pickling myself

>>> pndf = NamedDataFrame.read_pickle(mypath)
>>> pndf.name
AttributeError: 'NamedDataFrame' object has no attribute 'name'

问题根源 #

pandas的DataFrame类自定义了pickle的序列化逻辑（通过__reduce__方法），默认仅会保存DataFrame的核心数据（索引、列、值等结构信息），不会自动保存你添加的自定义属性。直接调用pickle.dump时，pandas的默认序列化流程会覆盖常规的实例属性保存逻辑，导致name属性未被写入pickle文件。

解决方案 #

通过重写__getstate__和__setstate__方法，手动将自定义属性纳入pickle的状态管理流程：

修改后的完整代码 #

import pandas as pd
import pickle

class NamedDataFrame(pd.DataFrame):
    '''
    a dataframe with a name
    '''
    def __init__(self, data=None, index=None, columns=None, dtype=None, copy=False, name: str = None):
        super().__init__(data, index, columns, dtype, copy)
        self.name = name
    
    def __getstate__(self):
        # 获取DataFrame的默认序列化状态
        state = super().__getstate__()
        # 将自定义属性加入状态字典
        state['name'] = self.name
        return state
    
    def __setstate__(self, state):
        # 先恢复DataFrame的核心状态
        super().__setstate__(state)
        # 恢复自定义属性
        self.name = state.get('name')
    
    def to_pickle(self, path, compression='infer', protocol=4):
        print("pickling myself")
        with open(path, 'wb') as f:
            pickle.dump(self, f, protocol)

    @classmethod
    def read_pickle(cls, path):
        with open(path, 'rb') as f:
            return pickle.load(f)

验证效果 #

>>> ndf = NamedDataFrame(data=mydf,name='mytestname')
>>> ndf.to_pickle(mypath)
pickling myself

>>> pndf = NamedDataFrame.read_pickle(mypath)
>>> pndf.name
'mytestname'

简化方案：无需重写pickle方法 #

只要正确实现__getstate__和__setstate__，可以直接使用pandas原生的to_pickle和read_pickle方法，无需自定义：

# 移除自定义的to_pickle和read_pickle方法后
pd.to_pickle(ndf, mypath)
pndf = pd.read_pickle(mypath)
print(pndf.name)  # 输出 'mytestname'

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