如何在PySpark Pipeline中使用XGBoost模型？

你遇到的问题很常见——直接用xgboost.XGBClassifier（sklearn版本）没法放进PySpark Pipeline，因为PySpark的Pipeline要求每个stage必须是Spark MLlib体系下的Estimator或Transformer，而sklearn模型并不属于这个范畴。不过别担心，XGBoost官方提供了专门适配Spark的版本，完全可以无缝集成到Pipeline里。

下面是具体的实现步骤和示例代码：

1. 先安装适配Spark的XGBoost包 #

确保你安装了支持Spark的XGBoost库，注意要和你的Spark版本兼容（比如Spark 3.3.x对应xgboost-spark 1.7+）：

pip install xgboost-spark pyspark

2. 导入正确的类 #

不要用sklearn的XGBClassifier，而是导入Spark版的：

from pyspark.sql import SparkSession
from pyspark.ml import Pipeline
from pyspark.ml.feature import VectorAssembler, StringIndexer
from xgboost.spark import XGBoostClassifier

3. 构建完整的Pipeline示例 #

这里以经典的分类任务为例，从数据准备到Pipeline训练、预测：

# 初始化SparkSession
spark = SparkSession.builder.appName("XGBoostPipelineExample").getOrCreate()

# 准备示例数据（替换成你的数据集）
data = spark.createDataFrame(
    [
        (1.0, 2.0, 0),
        (2.0, 3.0, 1),
        (3.0, 4.0, 0),
        (4.0, 5.0, 1),
    ],
    ["feature1", "feature2", "label"]
)

# 步骤1：把标签列转换成索引（如果标签是字符串的话需要这一步，数字标签可省略）
label_indexer = StringIndexer(inputCol="label", outputCol="indexed_label")

# 步骤2：把所有特征列合并成一个Vector类型的列（Spark XGBoost要求特征是Vector格式）
assembler = VectorAssembler(
    inputCols=["feature1", "feature2"],
    outputCol="features"
)

# 步骤3：初始化Spark版XGBoost分类器
xgb_classifier = XGBoostClassifier(
    featuresCol="features",
    labelCol="indexed_label",
    predictionCol="prediction",
    maxDepth=3,
    learningRate=0.1,
    nEstimators=10
)

# 构建Pipeline，按顺序添加所有stage
pipeline = Pipeline(stages=[label_indexer, assembler, xgb_classifier])

# 拆分训练集和测试集
train_data, test_data = data.randomSplit([0.8, 0.2], seed=42)

# 训练Pipeline模型
model = pipeline.fit(train_data)

# 用测试集预测
predictions = model.transform(test_data)

# 查看预测结果
predictions.select("features", "indexed_label", "prediction", "probability").show()

关键注意事项 #

• 版本兼容：一定要保证xgboost-spark的版本和你的PySpark版本匹配，否则会出现依赖错误。比如Spark 3.1.x对应xgboost-spark 1.5.x，Spark 3.3.x对应xgboost-spark 1.7.x+。
• 特征格式：Spark版XGBoost要求输入特征必须是Vector类型，所以必须用VectorAssembler把多列特征合并成一列。
• 参数差异：Spark版XGBoost的参数名称和sklearn版略有不同，比如sklearn里的n_estimators在Spark版里是nEstimators，max_depth是maxDepth，注意驼峰命名。
• 模型保存/加载：训练好的Pipeline模型可以直接用model.save()保存，之后用PipelineModel.load()加载，完全兼容Spark的模型持久化机制。

这样你就能像使用LogisticRegression一样，把XGBoost无缝集成到PySpark Pipeline里了！

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