PySpark 2.4.0含OneHotEncoder的PipelineModel迁移至3.0.0报错求助

问题原因 #

PySpark 3.0对OneHotEncoder组件做了重构：原2.4版本中已被标记为弃用的OneHotEncoder（Estimator类型）被移除，原OneHotEncoderEstimator重命名为OneHotEncoder，且PipelineModel要求所有stage必须是Transformer类型。你在2.4中用旧版OneHotEncoder训练的PipelineModel，保存的元数据中stage的类名仍为旧的org.apache.spark.ml.feature.OneHotEncoder（Estimator），在3.0加载时会被判定为非Transformer，从而触发报错。

解决方案（需在PySpark 2.4.0环境中操作） #

方式一：代码修改元数据并重保存模型 #

通过代码加载模型，替换旧stage的元数据为兼容3.0的Transformer类型，再重新保存：

from pyspark.ml import PipelineModel
from pyspark.ml.feature import OneHotEncoderModel

# 加载原模型
pipeline_model = PipelineModel.load("xxx/test_ohe")

modified_stages = []
for stage in pipeline_model.stages:
    if isinstance(stage, OneHotEncoderModel):
        # 修正元数据中的类名
        meta = stage._metadata
        meta["class"] = "org.apache.spark.ml.feature.OneHotEncoderModel"
        
        # 创建新的OneHotEncoderModel实例并替换元数据
        new_stage = OneHotEncoderModel(
            inputCol=stage.getInputCol(),
            outputCol=stage.getOutputCol(),
            dropLast=stage.getDropLast(),
            categories=stage.categories
        )
        new_stage._metadata = meta
        modified_stages.append(new_stage)
    else:
        modified_stages.append(stage)

# 生成新的PipelineModel并保存
new_pipeline_model = PipelineModel(stages=modified_stages)
new_pipeline_model.save("xxx/test_ohe_fixed")

方式二：手动修改模型元数据文件 #

如果无法运行代码，可直接修改模型保存目录的元数据文件：

• 进入模型保存目录的stages子文件夹，找到对应OneHotEncoder的stage目录（如OneHotEncoder_12f259fd25f8）
• 打开该目录下的metadata/part-00000文件
• 将文件中的"class": "org.apache.spark.ml.feature.OneHotEncoder"修改为"class": "org.apache.spark.ml.feature.OneHotEncoderModel"
• 保存文件后，即可在PySpark 3.0中正常加载模型

验证 #

在PySpark 3.0环境中加载修改后的test_ohe_fixed模型，执行transform操作测试是否正常运行。

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