本地加载Databricks用pickle/joblib保存的文件时遇pyspark.sql.metrics缺失错误

问题 #

在Databricks中执行以下代码读取并保存数据：

import pandas as pd
import joblib
query = 'select * from table a'
df = spark.sql(query)
df = df.toPandas()

df.to_pickle('df.pickle')
joblib.dump(df, 'df.joblib')

随后在本地PC尝试加载保存的文件：

import joblib
import pandas as pd
df = joblib.load('df.joblib')

触发如下错误：

ModuleNotFoundError: No module named 'pyspark.sql.metrics'
Cell In[8], line 1
----> 1 df = joblib.load('Data/df.joblib')
Hide Traceback
File ~\myenv\Lib\site-packages\joblib\numpy_pickle.py:658, in load(filename, mmap_mode)
    652             if isinstance(fobj, str):
    653                 # if the returned file object is a string, this means we
    654                 # try to load a pickle file generated with an version of
    655                 # Joblib so we load it with joblib compatibility function.
    656                 return load_compatibility(fobj)
--> 658             obj = _unpickle(fobj, filename, mmap_mode)
    659 return obj

File ~\myenv\Lib\site-packages\joblib\numpy_pickle.py:577, in _unpickle(fobj, filename, mmap_mode)
    575 obj = None
    576 try:
--> 577     obj = unpickler.load()
    578     if unpickler.compat_mode:
    579         warnings.warn("The file '%s' has been generated with a "
    580                       "joblib version less than 0.10. "
    581                       "Please regenerate this pickle file."
    582                       % filename,
    583                       DeprecationWarning, stacklevel=3)

File ~\AppData\Local\Programs\Python\Python311\Lib\pickle.py:1213, in _Unpickler.load(self)
   1211             raise EOFError
   1212         assert isinstance(key, bytes_types)
--> 1213         dispatch[key[0]](self)
   1214 except _Stop as stopinst:
   1215     return stopinst.value

File ~\AppData\Local\Programs\Python\Python311\Lib\pickle.py:1538, in _Unpickler.load_stack_global(self)
   1536 if type(name) is not str or type(module) is not str:
   1537     raise UnpicklingError("STACK_GLOBAL requires str")
--> 1538 self.append(self.find_class(module, name))

File ~\AppData\Local\Programs\Python\Python311\Lib\pickle.py:1580, in _Unpickler.find_class(self, module, name)
   1578     elif module in _compat_pickle.IMPORT_MAPPING:
   1579         module = _compat_pickle.IMPORT_MAPPING[module]
--> 1580 __import__(module, level=0)

解决办法 #

方案1：保存前清除Spark相关元数据 #

Spark DataFrame转Pandas后，部分列可能残留Spark相关类型或元数据，导致序列化时携带了Spark依赖对象。可以通过重新构造Pandas DataFrame彻底剥离这些依赖：

# 在Databricks中修改保存代码
import pandas as pd
import joblib
query = 'select * from table a'
df = spark.sql(query)
df = df.toPandas()

# 重新构造DataFrame，确保所有列都是纯Pandas类型
clean_df = pd.DataFrame(df.values, columns=df.columns)

clean_df.to_pickle('df.pickle')
joblib.dump(clean_df, 'df.joblib')

方案2：使用无Spark依赖的通用存储格式 #

避免依赖环境的序列化方式，改用CSV、Parquet等兼容性更好的格式：

# Databricks中保存为Parquet（推荐，体积小、读写效率高）
df.to_parquet('df.parquet')

# 本地加载
df = pd.read_parquet('df.parquet')

# 或者保存为CSV
df.to_csv('df.csv', index=False)
# 本地加载
df = pd.read_csv('df.csv')

方案3：本地安装对应版本PySpark（不推荐） #

如果坚持使用pickle/joblib加载，可在本地环境安装与Databricks一致版本的PySpark：

pip install pyspark==<Databricks中PySpark版本>

此方式会引入不必要的依赖，且版本不匹配仍可能出现问题，不建议优先使用。

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