pyspark中的模型预测

Question

我建立了一个xgboost模型（在python中）并将其保存为.pkl格式。现在，我想在查找.pkl文件后预测一个新的数据集。如何在Pyspark中做到这一点。

我的df-

id     score1    score2    score3    
1      1.346     .987       .199
2      .976      1.43       .999
3       .098     2.324      1.33
4      1.33      .876       1.97

预测后的最终数据帧应为-

id     score1    score2    score3    score
1      1.346     .987       .199      1.876
2      .976      1.43       .999      2.876
3       .098     2.324      1.33      3.348
4      1.33      .876       1.97      2.98

我已经在pyspark中完成了以下代码-

df= spark.read.parquet('s3a://XXX/...../AC_data_cleaning_v1/')

model_path = SparkFiles.get("score.pkl")
loaded_model = pkl.load(open(model_path, 'rb'))

feature = xgb.DMatrix(df)
class_probs = loaded_model.predict(feature)

yield (int(i[0]), float(class_probs))