如何pickle包含Spark DataFrame属性的自定义Python对象

Spark DataFrame本身是JVM侧的执行计划+数据引用，携带线程锁、SparkContext绑定的运行时状态，无法直接被原生pickle序列化。要实现带DataFrame属性的自定义类单文件序列化，最稳妥的方式是自定义类的pickle序列化钩子，将类的所有状态打包为单份可序列化的字节流，对外保持和普通pickle完全一致的使用体验。

具体实现 #

通过Python原生的__getstate__（序列化时自动触发）、__setstate__（反序列化时自动触发）魔术方法，将三个DataFrame统一写入临时目录后打包为单个zip字节流，全程不需要手动管理多个独立文件，最终生成的就是单个可加载的对象文件。这里用Parquet格式存储DataFrame，比saveAsPickleFile的跨版本兼容性更好。

import os
import tempfile
import shutil
import zipfile
import pickle
from pyspark.sql import SparkSession

class MyClass:
    def __init__(self, df_a, df_b, df_c):
        self.a = df_a
        self.b = df_b
        self.c = df_c

    def __getstate__(self):
        tmp_dir = tempfile.mkdtemp()
        try:
            # 将三个DataFrame写入临时目录下的固定子路径
            self.a.write.mode("overwrite").parquet(os.path.join(tmp_dir, "df_a"))
            self.b.write.mode("overwrite").parquet(os.path.join(tmp_dir, "df_b"))
            self.c.write.mode("overwrite").parquet(os.path.join(tmp_dir, "df_c"))
            
            # 打包整个临时目录为单个zip字节流
            zip_buffer = tempfile.SpooledTemporaryFile()
            with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
                for root, _, files in os.walk(tmp_dir):
                    for file in files:
                        file_path = os.path.join(root, file)
                        arcname = os.path.relpath(file_path, tmp_dir)
                        zf.write(file_path, arcname)
            zip_buffer.seek(0)
            return {"packed_data": zip_buffer.read()}
        finally:
            shutil.rmtree(tmp_dir)

    def __setstate__(self, state):
        # 反序列化前必须提前初始化激活的SparkSession
        spark = SparkSession.getActiveSession()
        if not spark:
            raise RuntimeError("请先初始化SparkSession再加载序列化对象")
        
        tmp_dir = tempfile.mkdtemp()
        try:
            # 解压zip包到临时目录
            zip_buffer = tempfile.SpooledTemporaryFile()
            zip_buffer.write(state["packed_data"])
            zip_buffer.seek(0)
            with zipfile.ZipFile(zip_buffer, "r") as zf:
                zf.extractall(tmp_dir)
            
            # 读取Parquet恢复DataFrame属性
            self.a = spark.read.parquet(os.path.join(tmp_dir, "df_a"))
            self.b = spark.read.parquet(os.path.join(tmp_dir, "df_b"))
            self.c = spark.read.parquet(os.path.join(tmp_dir, "df_c"))
        finally:
            shutil.rmtree(tmp_dir)

使用方式 #

和普通Python对象的pickle逻辑完全一致，最终只生成单个文件：

# 保存实例到单个pkl文件
obj = MyClass(df_a, df_b, df_c)
with open("my_saved_obj.pkl", "wb") as f:
    pickle.dump(obj, f)

# 加载实例
spark = SparkSession.builder.appName("load_demo").getOrCreate()
with open("my_saved_obj.pkl", "rb") as f:
    loaded_obj = pickle.load(f)

# 直接使用恢复的对象
loaded_obj.a.show(5)

优化提示 #

• 如果DataFrame数据量极大，不想把全量数据打包进序列化文件，可以修改__getstate__逻辑，只存储三个DataFrame对应的数据源元信息（比如Hive表名、对象存储路径、过滤条件、分区参数），反序列化时直接通过SparkSession读取对应数据源即可，最终生成的序列化文件体积会非常小。
• 不要尝试直接pickle DataFrame的内部JVM对象引用，脱离原SparkContext/JVM进程后这类引用会直接失效，跨进程、跨节点加载必然报错。

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