<jupyter_start><jupyter_text>`spark-xgboost` PySpark wrapper for XGBoost4J-SparkThis notebook is a brief example of how to use the XGBoost4J-Spark library from python. In it, we train a model to predict earning potential from the 'Adult' dataset by putting the model into a standard SparkML pipeline.Prerequisites:- Databricks cluster with Databricks Runtime for ML >= 6.6_or_- Spark cluster with Spark >= 2.4.5, PySpark >= 2.4.5, XGBoost >= 0.9, MLflow >= 1.8<jupyter_code>from sparkxgb import XGBoostClassifier, XGBoostRegressor
from pprint import PrettyPrinter

from pyspark.sql.types import StringType

from pyspark.ml import Pipeline
from pyspark.ml.feature import StringIndexer, VectorAssembler
from pyspark.ml.evaluation import MulticlassClassificationEvaluator, BinaryClassificationEvaluator
from pyspark.ml.tuning import ParamGridBuilder, CrossValidator
pp = PrettyPrinter()
print(pp)
import os
def get_running_jars(workdir):
    import glob
    return glob.glob("{}/*jar".format(workdir))

workspace=os.path.abspath(os.path.dirname(os.path.realpath('__file__')))
print(workspace)
spark_jars = get_running_jars("{}/../jars/".format(workspace))
print(spark_jars)

#os.environ['PYSPARK_SUBMIT_ARGS'] = "--jars {} pyspark-shell".format(" ,".join(spark_jars))
from pyspark.context import SparkContext

#sc = SparkSession.builder.config("spark.jars",",".join(spark_jars)).getOrCreate()
print(os.environ['PYSPARK_SUBMIT_ARGS'])
spark.sparkContext.getConf().getAll()

col_names = [
  "age", "workclass", "fnlwgt",
  "education", "education-num",
  "marital-status", "occupation",
  "relationship", "race", "sex",
  "capital-gain", "capital-loss",
  "hours-per-week", "native-country",
  "label"
]

# https://archive.ics.uci.edu/ml/machine-learning-databases/adult/
train_sdf, test_sdf = (
  spark.read.csv(
    path="../sparkxgb/tests/data/adult.data",
    inferSchema=True  
  )
  .toDF(*col_names)
  .repartition(200)
  .randomSplit([0.8, 0.2])
)
string_columns = [fld.name for fld in train_sdf.schema.fields if isinstance(fld.dataType, StringType)]
string_col_replacements = [fld + "_ix" for fld in string_columns]
string_column_map=list(zip(string_columns, string_col_replacements))
target = string_col_replacements[-1]
predictors = [fld.name for fld in train_sdf.schema.fields if not isinstance(fld.dataType, StringType)] + string_col_replacements[:-1]
pp.pprint(
  dict(
    string_column_map=string_column_map,
    target_variable=target,
    predictor_variables=predictors
  )
)
si = [StringIndexer(inputCol=fld[0], outputCol=fld[1]) for fld in string_column_map]
va = VectorAssembler(inputCols=predictors, outputCol="features")
pipeline = Pipeline(stages=[*si, va])
fitted_pipeline = pipeline.fit(train_sdf.union(test_sdf))
train_sdf_prepared = fitted_pipeline.transform(train_sdf)
train_sdf_prepared.cache()
train_sdf_prepared.count()
test_sdf_prepared = fitted_pipeline.transform(test_sdf)
test_sdf_prepared.cache()
test_sdf_prepared.count()
xgbParams = dict(
  eta=0.1,
  maxDepth=2,
  missing=0.0,
  objective="binary:logistic",
  numRound=5,
  numWorkers=2
)

xgb = (
  XGBoostClassifier(**xgbParams)
  .setFeaturesCol("features")
  .setLabelCol("label_ix")
)

bce = BinaryClassificationEvaluator(
  rawPredictionCol="rawPrediction",
  labelCol="label_ix"
)
print(xgb)
print(bce)
param_grid = (
  ParamGridBuilder()
  .addGrid(xgb.eta, [1e-1, 1e-2, 1e-3])
  .addGrid(xgb.maxDepth, [2, 4, 8])
  .build()
)

cv = CrossValidator(
  estimator=xgb,
  estimatorParamMaps=param_grid,
  evaluator=bce,#mce,
  numFolds=5
)
print(cv)
import mlflow
import mlflow.spark

spark_model_name = "best_model_spark"

with mlflow.start_run():
  model = cv.fit(train_sdf_prepared)
  best_params = dict(
    eta_best=model.bestModel.getEta(),
    maxDepth_best=model.bestModel.getMaxDepth()
  )
  mlflow.log_params(best_params)
  
  mlflow.spark.log_model(fitted_pipeline, "featuriser")
  mlflow.spark.log_model(model, spark_model_name)

  metrics = dict(
    roc_test=bce.evaluate(model.transform(test_sdf_prepared))
  )
  mlflow.log_metrics(metrics)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12809815950920245, 'fraction_numerical': 0.010306748466257669, 'mean_word_length': 5.931972789115647, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 10, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '384871', 'license': 'Apache-2.0', 'n_tokens_mistral': 1562, 'n_tokens_neox': 1509, 'n_words': 264, 'path': 'examples/spark-xgboost_adultdataset.ipynb', 'repo_name': 'westfly/spark-xgboost', 'size': '4075'}