Azure Databricks中SparkXGBRegressorModel无法写入Feature Store求助

问题出在你用了mlflow.xgboost作为flavor，而SparkXGBRegressorModel属于PySpark ML模型，并非原生XGBoost模型，所以不兼容mlflow.xgboost的序列化逻辑（该flavor会调用模型的save_model方法，但Spark模型没有这个方法）。

方案一：使用PySpark ML flavor（推荐） #

直接改用mlflow.pyspark.ml作为flavor，适配Spark模型的存储逻辑：

fs.log_model(
    model=xgb_model,
    artifact_path=f"{model_name}_artifact_path",
    flavor=mlflow.pyspark.ml,
    training_set=training_set,
    registered_model_name=model_name
)

这个方案会完整保留SparkXGBRegressorModel的所有特性，后续可以直接在Spark环境中通过Feature Store加载模型，用Spark UDF或Pipeline进行批量预测，和Databricks的生态集成更顺畅。

方案二：提取原生XGBoost模型记录 #

如果后续需要在非Spark环境使用原生XGBoost模型，可以先从SparkXGBRegressorModel中提取出原生的Booster对象，再用mlflow.xgboost flavor记录：

# 提取原生XGBoost Booster模型
native_xgb_booster = xgb_model._xgb_regressor.get_booster()

# 记录到Feature Store
fs.log_model(
    model=native_xgb_booster,
    artifact_path=f"{model_name}_native_artifact_path",
    flavor=mlflow.xgboost,
    training_set=training_set,
    registered_model_name=f"{model_name}_native"
)

注意：这种方式会丢失Spark模型的封装，无法直接在Spark Pipeline中使用，仅适合原生XGBoost部署场景。

内容的提问来源于stack exchange，提问作者Φιλιππος Πριοβολος