如何保存并恢复带自定义属性的pandas DataFrame

别担心，继承pandas.DataFrame绝对是解决这个问题的绝佳思路！不仅能完美保留你加的自定义元属性，还能解决CSV/JSON丢失类型、索引的痛点。下面给你几个实用的具体方案：

Pickle本身就能完整保存DataFrame的结构、类型和索引，我们只需要把自定义元属性和DataFrame一起打包存储，加载时再恢复即可。直接上代码：

import pandas as pd
import pickle

class MetaDataFrame(pd.DataFrame):
    def __init__(self, *args, **kwargs):
        # 先把自定义元属性从参数里拎出来，剩下的交给父类处理
        self.meta = kwargs.pop('meta', {})
        super().__init__(*args, **kwargs)
    
    def save(self, filepath):
        # 把DataFrame和元属性打包成字典保存
        with open(filepath, 'wb') as f:
            pickle.dump({
                'data': self,
                'meta': self.meta
            }, f)
    
    @classmethod
    def load(cls, filepath):
        # 加载时先读取打包内容，再恢复成自定义类实例
        with open(filepath, 'rb') as f:
            saved_content = pickle.load(f)
            # 把加载的DataFrame转成我们的自定义类
            loaded_df = cls(saved_content['data'])
            # 绑定元属性
            loaded_df.meta = saved_content['meta']
            return loaded_df

用法示例 #

# 创建带元属性的自定义DataFrame
df = MetaDataFrame(
    {'user_id': [101, 102, 103], 'score': [85, 92, 78]},
    meta={'dataset': 'user_scores', 'created_at': '2024-05-20', 'version': 2.1}
)

# 保存到文件
df.save('user_data.pkl')

# 加载文件
loaded_df = MetaDataFrame.load('user_data.pkl')

# 验证元属性是否保留
print(loaded_df.meta)
# 输出: {'dataset': 'user_scores', 'created_at': '2024-05-20', 'version': 2.1}

# 验证数据完整性
print(loaded_df.equals(df))
# 输出: True

注意：Pickle只适合内部数据交换，不要加载来自不可信来源的pickle文件，存在安全风险。

如果必须用CSV/JSON这类可读格式，可以把元属性单独存到一个JSON文件里，加载时一起读取恢复。这样既保证了数据的可读性，又能保留元信息：

import pandas as pd
import json

class MetaDataFrame(pd.DataFrame):
    def __init__(self, *args, **kwargs):
        self.meta = kwargs.pop('meta', {})
        super().__init__(*args, **kwargs)
    
    def save_csv(self, data_filepath, meta_filepath=None):
        # 保存数据到CSV（按需决定是否保留索引）
        self.to_csv(data_filepath, index=True)
        # 默认元数据文件和数据文件同名，后缀改为_meta.json
        if not meta_filepath:
            meta_filepath = '.'.join(data_filepath.split('.')[:-1]) + '_meta.json'
        # 把元属性写入JSON文件
        with open(meta_filepath, 'w', encoding='utf-8') as f:
            json.dump(self.meta, f, indent=2)
    
    @classmethod
    def load_csv(cls, data_filepath, meta_filepath=None):
        # 加载CSV数据（记得按需解析日期等特殊类型）
        df = pd.read_csv(data_filepath, index_col=0)
        # 加载对应元数据
        if not meta_filepath:
            meta_filepath = '.'.join(data_filepath.split('.')[:-1]) + '_meta.json'
        with open(meta_filepath, 'r', encoding='utf-8') as f:
            meta = json.load(f)
        # 转为自定义类实例并绑定元属性
        return cls(df, meta=meta)

补充优化 #

CSV本身会丢失一些复杂类型（比如datetime），你可以在元属性里记录列的类型信息，加载时自动转换：

# 保存时记录列类型
df.meta['dtypes'] = {col: str(dtype) for col, dtype in df.dtypes.items()}

# 加载时根据记录的类型转换
for col, dtype in loaded_df.meta['dtypes'].items():
    loaded_df[col] = loaded_df[col].astype(eval(dtype))

如果是处理大数据场景，Parquet是更好的选择——它不仅能完美保留DataFrame的类型、索引，还原生支持自定义元数据，体积更小、读取更快：

import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq

class MetaDataFrame(pd.DataFrame):
    def __init__(self, *args, **kwargs):
        self.meta = kwargs.pop('meta', {})
        super().__init__(*args, **kwargs)
    
    def save_parquet(self, filepath):
        # 转成pyarrow Table
        table = pa.Table.from_pandas(self)
        # 把自定义元属性编码后加入表的元数据
        encoded_meta = {k: v.encode('utf-8') for k, v in self.meta.items()}
        updated_metadata = {**table.schema.metadata, **encoded_meta}
        table = table.replace_metadata(updated_metadata)
        # 写入Parquet文件
        pq.write_table(table, filepath)
    
    @classmethod
    def load_parquet(cls, filepath):
        # 读取Parquet表
        table = pq.read_table(filepath)
        # 提取自定义元属性并解码
        meta = {}
        for key, value in table.schema.metadata.items():
            # 过滤掉pandas自带的元数据
            if not key.startswith(b'pandas'):
                meta[key.decode('utf-8')] = value.decode('utf-8')
        # 转回DataFrame并转为自定义类
        df = table.to_pandas()
        return cls(df, meta=meta)

如果你的自定义DataFrame需要复制操作，记得重写copy方法，确保元属性也被复制：

def copy(self, deep=True):
    copied_df = super().copy(deep=deep)
    # 深拷贝元属性，避免原对象和副本共享引用
    copied_meta = self.meta.copy() if deep else self.meta
    return self.__class__(copied_df, meta=copied_meta)

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