Python中含Polars DataFrame的字典JSON序列化与反序列化问题

问题描述 #

我是Python中Polars的新手，尝试保存部分值为pl.DataFrame的字典到JSON格式时遇到问题。初始代码因DataFrame无法JSON序列化报错，调整后可完成序列化，但无法将反序列化后的内容恢复为pl.DataFrame。

初始报错代码：

import polars as pl
import json

df = pl.DataFrame({'x': 1,
                   'y': 2,
                   'z': 3})

data = dict({'K01': ['001', '002', '003'], 
             'K02': {'K11': '1001', 
                     'K12': '1002', 
                     'K13': df
                     }})

data_jsn = json.dumps(data, indent = 4)
# TypeError: Object of type DataFrame is not JSON serializable

调整后的代码（可序列化但无法恢复）：

# write df to json
df_jsn = df.write_json()

# then put df_json into the dict in place of df
data_jsn = dict({'K01': ['001', '002', '003'], 
                 'K02': {'K11': '1001', 
                         'K12': '1002', 
                         'K13': df_jsn
                         }})

data_jsn_ser = json.dumps(data_jsn, indent = 4)
data_jsn_deser = json.loads(data_jsn_ser)

print(data_jsn_deser['K02']['K13'])
# 输出：{"columns":[{"name":"x","datatype":"Int64","bit_settings":"","values":[1]},{"name":"y","datatype":"Int64","bit_settings":"","values":[2]},{"name":"z","datatype":"Int64","bit_settings":"","values":[3]}]}

解答 #

1. 现有方法的挽救方案

你的调整方案中，data_jsn_deser['K02']['K13']是一个被JSON转义后的字符串，只需两步就能恢复为DataFrame：

# 取出转义后的DataFrame JSON字符串
df_json_str = data_jsn_deser['K02']['K13']
# 方式1：解析为字典后转DataFrame
df_dict = json.loads(df_json_str)
restored_df = pl.DataFrame.from_dict(df_dict)

# 方式2：直接用Polars读取JSON字符串
restored_df = pl.read_json(df_json_str)

执行后即可得到与原始df完全一致的Polars DataFrame。

2. 更优的实现方式

手动替换DataFrame为JSON字符串的方式不够高效，推荐用自定义JSON编码器/解码器自动处理Polars对象：

自定义编码器与解码器

import polars as pl
import json

class PolarsJSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, pl.DataFrame):
            # 标记类型并保存DataFrame的字典结构
            return {"__polars_df__": True, "data": obj.to_dict(as_series=False)}
        return super().default(obj)

def polars_decoder(obj):
    if obj.get("__polars_df__"):
        # 识别标记并恢复DataFrame
        return pl.DataFrame(obj["data"])
    return obj

# 序列化整个字典
data = {
    'K01': ['001', '002', '003'], 
    'K02': {'K11': '1001', 'K12': '1002', 'K13': df}
}
data_jsn_ser = json.dumps(data, cls=PolarsJSONEncoder, indent=4)

# 反序列化直接得到带DataFrame的字典
data_jsn_deser = json.loads(data_jsn_ser, object_hook=polars_decoder)

# 验证结果
print(type(data_jsn_deser['K02']['K13']))  # <class 'polars.dataframe.frame.DataFrame'>

其他可选方案

如果不需要严格的JSON格式，也可以选择：

• Pickle序列化：直接保存整个Python对象，操作简单但存在安全风险，仅用于可信环境：

  import pickle
  # 序列化
  with open("data.pkl", "wb") as f:
      pickle.dump(data, f)
  # 反序列化
  with open("data.pkl", "rb") as f:
      restored_data = pickle.load(f)
  

• Parquet格式：将字典中的DataFrame单独保存为Parquet文件，在字典中记录文件路径，反序列化时读取文件恢复，适合大数据量场景。

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