Databricks 中的 XGBoost Python

XGBoost in Databricks with Python

最近我一直在研究 Mlib Databricks 集群,根据文档发现 XGBoost 可用于我的集群版本 (5.1)。这个集群是 运行 Python 2.

我感觉 XGBoost4J 仅适用于 Scala 和 Java。 所以我的问题是:如何在不丢失分发功能的情况下将 xgboost 模块导入此环境?

下面是我的代码示例

from pyspark.ml import Pipeline
from pyspark.ml.feature import StringIndexer
from pyspark.ml.feature import VectorAssembler
import xgboost as xgb # Throws error because module is not installed and it should

# Transform class to classIndex to make xgboost happy
stringIndexer = StringIndexer(inputCol="species", outputCol="species_index").fit(newInput)
labelTransformed = stringIndexer.transform(newInput).drop("species")

# Compose feature columns as vectors
vectorCols = ["sepal_length", "sepal_width", "petal_length", "petal_width", "species_index"]
vectorAssembler = VectorAssembler(inputCols=vectorCols, outputCol="features")
xgbInput = vectorAssembler.transform(labelTransformed).select("features", "species_index")

您可以尝试使用 spark-sklearn 分发 python 或 scikit-learn 版本的 xgboost,但该分发与 xgboost4j 分发不同。听说databricks上xgboost4j的pyspark api要来了,敬请期待。

相关pull request,顺便可以找到here