Spark运行Xgboost且保存为PMML

重新修正

软件版本：
Spark版本2.1.3
xgboost4j-spark 0.81

---

maven依赖

  <properties>
        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
        <spark.version>2.1.3</spark.version>
        <hadoop.version>2.6.4</hadoop.version>
        <fastjson.version>1.2.47</fastjson.version>
    </properties>

    <dependencies>

        <dependency>
            <groupId>org.apache.spark</groupId>
            <artifactId>spark-core_2.11</artifactId>
            <version>${spark.version}</version>
        </dependency>

        <dependency>
            <groupId>org.apache.hadoop</groupId>
            <artifactId>hadoop-common</artifactId>
            <version>${hadoop.version}</version>
        </dependency>

        <dependency>
            <groupId>org.apache.spark</groupId>
            <artifactId>spark-sql_2.11</artifactId>
            <version>${spark.version}</version>
        </dependency>

        <dependency>
            <groupId>org.apache.spark</groupId>
            <artifactId>spark-hive_2.11</artifactId>
            <version>${spark.version}</version>
        </dependency>

        <dependency>
            <groupId>org.apache.spark</groupId>
            <artifactId>spark-mllib_2.11</artifactId>
            <version>${spark.version}</version>
            <exclusions>
                <exclusion>
                    <groupId>org.jpmml</groupId>
                    <artifactId>pmml-model</artifactId>
                </exclusion>
            </exclusions>
        </dependency>

        <dependency>
            <groupId>org.jpmml</groupId>
            <artifactId>jpmml-sparkml</artifactId>
            <version>1.2.14</version>
        </dependency>

        <dependency>
            <groupId>ml.dmlc</groupId>
            <artifactId>xgboost4j-spark</artifactId>
            <version>0.81</version>
        </dependency>

        <dependency>
            <groupId>org.jpmml</groupId>
            <artifactId>jpmml-xgboost</artifactId>
            <version>1.3.3</version>
        </dependency>
    </dependencies>

说明，如果想保存将XGBoost模型保存为PMML格式，
spark-ml必须exclude org.jpmml:pmml-model.
ml.dmlc:xgboost4j-spark:0.81是适配spark2.3的
如果想在spark2.1.3运行，在项目下新建org.apache.spark包，
新建SparkParallelismTracker，
将sc.removeListener(listener)修改sc.listenerBus.removeListener(listener)

/*
 Copyright (c) 2014 by Contributors

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

 http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 */

package org.apache.spark

import java.net.URL

import org.apache.commons.logging.LogFactory

import org.apache.spark.scheduler.{SparkListener, SparkListenerExecutorRemoved, SparkListenerTaskEnd}
import org.codehaus.jackson.map.ObjectMapper
import scala.collection.JavaConverters._
import scala.concurrent.ExecutionContext.Implicits.global
import scala.concurrent.duration._
import scala.concurrent.{Await, Future, TimeoutException}
import scala.util.control.ControlThrowable

/**
  * A tracker that ensures enough number of executor cores are alive.
  * Throws an exception when the number of alive cores is less than nWorkers.
  *
  * @param sc The SparkContext object
  * @param timeout The maximum time to wait for enough number of workers.
  * @param numWorkers nWorkers used in an XGBoost Job
  */
class SparkParallelismTracker(
                               val sc: SparkContext,
                               timeout: Long,
                               numWorkers: Int) {

  private[this] val requestedCores = numWorkers * sc.conf.getInt("spark.task.cpus", 1)
  private[this] val mapper = new ObjectMapper()
  private[this] val logger = LogFactory.getLog("XGBoostSpark")
  private[this] val url = sc.uiWebUrl match {
    case Some(baseUrl) => new URL(s"$baseUrl/api/v1/applications/${sc.applicationId}/executors")
    case _ => null
  }

  private[this] def numAliveCores: Int = {
    try {
      if (url != null) {
        mapper.readTree(url).findValues("totalCores").asScala.map(_.asInt).sum
      } else {
        Int.MaxValue
      }
    } catch {
      case ex: Throwable =>
        logger.warn(s"Unable to read total number of alive cores from REST API." +
          s"Health Check will be ignored.")
        ex.printStackTrace()
        Int.MaxValue
    }
  }

  private[this] def waitForCondition(
                                      condition: => Boolean,
                                      timeout: Long,
                                      checkInterval: Long = 100L) = {
    val monitor = Future {
      while (!condition) {
        Thread.sleep(checkInterval)
      }
    }
    Await.ready(monitor, timeout.millis)
  }

  private[this] def safeExecute[T](body: => T): T = {
    val listener = new TaskFailedListener
    sc.addSparkListener(listener)
    try {
      body
    } finally {
      //注销地方
//      sc.removeListener(listener)
      sc.listenerBus.removeListener(listener)
    }
  }

  /**
    * Execute a blocking function call with two checks on enough nWorkers:
    *  - Before the function starts, wait until there are enough executor cores.
    *  - During the execution, throws an exception if there is any executor lost.
    *
    * @param body A blocking function call
    * @tparam T Return type
    * @return The return of body
    */
  def execute[T](body: => T): T = {
    if (timeout <= 0) {
      logger.info("starting training without setting timeout for waiting for resources")
      body
    } else {
      try {
        logger.info(s"starting training with timeout set as $timeout ms for waiting for resources")
        waitForCondition(numAliveCores >= requestedCores, timeout)
      } catch {
        case _: TimeoutException =>
          throw new IllegalStateException(s"Unable to get $requestedCores workers for" +
            s" XGBoost training")
      }
      safeExecute(body)
    }
  }
}

private class ErrorInXGBoostTraining(msg: String) extends ControlThrowable {
  override def toString: String = s"ErrorInXGBoostTraining: $msg"
}

private[spark] class TaskFailedListener extends SparkListener {
  override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = {
    taskEnd.reason match {
      case taskEnd: SparkListenerTaskEnd =>
        if (taskEnd.reason.isInstanceOf[TaskFailedReason]) {
          throw new ErrorInXGBoostTraining(s"TaskFailed during XGBoost Training: " +
            s"${taskEnd.reason}")
        }
      case executorRemoved: SparkListenerExecutorRemoved =>
        throw new ErrorInXGBoostTraining(s"Executor lost during XGBoost Training: " +
          s"${executorRemoved.reason}")
      case _ =>
    }
  }
}

保存PMML，需要新建org.jpmml.sparkml.xgboost包
新建BoosterUtil、HasXGBoostOptions、XGBoostClassificationModelConverter、XGBoostRegressionModelConverter四个类。参考：jpmml-sparkml-xgboost
编译成依赖包也行，我这里直接放在代码里了。

/*
 * Copyright (c) 2017 Villu Ruusmann
 *
 * This file is part of JPMML-SparkML
 *
 * JPMML-SparkML is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * JPMML-SparkML is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with JPMML-SparkML.  If not, see <http://www.gnu.org/licenses/>.
 */
package org.jpmml.sparkml.xgboost;

import ml.dmlc.xgboost4j.scala.Booster;
import org.dmg.pmml.DataType;
import org.dmg.pmml.mining.MiningModel;
import org.jpmml.converter.BinaryFeature;
import org.jpmml.converter.ContinuousFeature;
import org.jpmml.converter.Feature;
import org.jpmml.converter.Schema;
import org.jpmml.sparkml.ModelConverter;
import org.jpmml.xgboost.Learner;
import org.jpmml.xgboost.XGBoostUtil;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.function.Function;

public class BoosterUtil {

	private BoosterUtil(){
	}

	static
	public <C extends ModelConverter<?> & HasXGBoostOptions> MiningModel encodeBooster(C converter, Booster booster, Schema schema){
		byte[] bytes = booster.toByteArray();

		Learner learner;

		try(InputStream is = new ByteArrayInputStream(bytes)){
			learner = XGBoostUtil.loadLearner(is);
		} catch(IOException ioe){
			throw new RuntimeException(ioe);
		}

		Function<Feature, Feature> function = new Function<Feature, Feature>(){

			@Override
			public Feature apply(Feature feature){

				if(feature instanceof BinaryFeature){
					BinaryFeature binaryFeature = (BinaryFeature)feature;

					return binaryFeature;
				} else

				{
					ContinuousFeature continuousFeature = feature.toContinuousFeature(DataType.FLOAT);

					return continuousFeature;
				}
			}
		};

		Map<String, Object> options = new LinkedHashMap<>();
		options.put(org.jpmml.xgboost.HasXGBoostOptions.OPTION_COMPACT, converter.getOption(HasXGBoostOptions.OPTION_COMPACT, false));
		options.put(org.jpmml.xgboost.HasXGBoostOptions.OPTION_NTREE_LIMIT, converter.getOption(HasXGBoostOptions.OPTION_NTREE_LIMIT, null));

		Schema xgbSchema = schema.toTransformedSchema(function);

		return learner.encodeMiningModel(options, xgbSchema);
	}
}

/*
 * Copyright (c) 2018 Villu Ruusmann
 *
 * This file is part of JPMML-SparkML
 *
 * JPMML-SparkML is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * JPMML-SparkML is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with JPMML-SparkML.  If not, see <http://www.gnu.org/licenses/>.
 */
package org.jpmml.sparkml.xgboost;

import org.jpmml.converter.HasOptions;


public interface HasXGBoostOptions extends HasOptions {

	String OPTION_COMPACT = "compact";

	String OPTION_NTREE_LIMIT = "ntree_limit";
}

/*
 * Copyright (c) 2017 Villu Ruusmann
 *
 * This file is part of JPMML-SparkML
 *
 * JPMML-SparkML is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * JPMML-SparkML is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with JPMML-SparkML.  If not, see <http://www.gnu.org/licenses/>.
 */
package org.jpmml.sparkml.xgboost;

import ml.dmlc.xgboost4j.scala.Booster;
import ml.dmlc.xgboost4j.scala.spark.XGBoostClassificationModel;
import org.dmg.pmml.mining.MiningModel;
import org.jpmml.converter.Schema;
import org.jpmml.sparkml.ClassificationModelConverter;


public class XGBoostClassificationModelConverter extends ClassificationModelConverter/*<XGBoostClassificationModel>*/ implements HasXGBoostOptions {

	public XGBoostClassificationModelConverter(XGBoostClassificationModel model){
		super(model);
	}

	@Override
	public MiningModel encodeModel(Schema schema){
		XGBoostClassificationModel model = (XGBoostClassificationModel)getTransformer();

		Booster booster = model.nativeBooster();

		return BoosterUtil.encodeBooster(this, booster, schema);
	}
}

/*
 * Copyright (c) 2017 Villu Ruusmann
 *
 * This file is part of JPMML-SparkML
 *
 * JPMML-SparkML is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * JPMML-SparkML is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with JPMML-SparkML.  If not, see <http://www.gnu.org/licenses/>.
 */
package org.jpmml.sparkml.xgboost;

import ml.dmlc.xgboost4j.scala.Booster;
import ml.dmlc.xgboost4j.scala.spark.XGBoostRegressionModel;
import org.dmg.pmml.mining.MiningModel;
import org.jpmml.converter.Schema;
import org.jpmml.sparkml.RegressionModelConverter;

public class XGBoostRegressionModelConverter extends RegressionModelConverter/*<XGBoostRegressionModel>*/ implements HasXGBoostOptions {

	public XGBoostRegressionModelConverter(XGBoostRegressionModel model){
		super(model);
	}

	@Override
	public MiningModel encodeModel(Schema schema){
		XGBoostRegressionModel model = (XGBoostRegressionModel)getTransformer();

		Booster booster = model.nativeBooster();

		return BoosterUtil.encodeBooster(this, booster, schema);
	}
}

需要在resource下新建META-INF目录，新建文件sparkml2pmml.properties
配置内容如下：

ml.dmlc.xgboost4j.scala.spark.XGBoostClassificationModel = org.jpmml.sparkml.xgboost.XGBoostClassificationModelConverter
ml.dmlc.xgboost4j.scala.spark.XGBoostRegressionModel = org.jpmml.sparkml.xgboost.XGBoostRegressionModelConverter

---

运行demo

package com.demo.ml;


import com.demo.base.HDFSFileSystem;
import com.demo.util.CommonUtil;
import ml.dmlc.xgboost4j.java.XGBoostError;
import ml.dmlc.xgboost4j.scala.spark.XGBoostClassificationModel;


import ml.dmlc.xgboost4j.scala.spark.XGBoostClassifier;
import org.apache.hadoop.fs.Path;
import org.apache.spark.ml.Pipeline;
import org.apache.spark.ml.PipelineModel;
import org.apache.spark.ml.PipelineStage;
import org.apache.spark.ml.evaluation.BinaryClassificationEvaluator;
import org.apache.spark.ml.feature.VectorAssembler;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.types.StructType;
import org.dmg.pmml.PMML;
import org.jpmml.model.PMMLUtil;
import org.jpmml.sparkml.PMMLBuilder;

import javax.xml.bind.JAXBException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

/**
 * @author allen
 * @date 2019/1/14.
 */
public class XGBoostLocalDemo {

	public static void main(String[] args) throws IOException, XGBoostError, JAXBException {

		SparkSession spark =SparkSession.builder().master("local[*]").getOrCreate();

		Dataset<Row> trainData = spark.read()
				.option("inferschema", "true")
				.option("header", "true")
				.option("encoding", "gbk")
				.csv("/Users/AllenBai/data/ml_offical_data.csv")
				.drop("userid,city,from_place,to_place,start_city_name,end_city_name,start_city_id,end_city_id,create_date".split(","));;

		trainData.printSchema();
		trainData.show(false);

		trainData.cache();
		String[] features=new String[]{
				"category", "future_day",
				"banner_min_time","banner_min_price",
				"page_train", "page_flight", "page_bus",
				"page_transfer",
				"start_end_distance", "total_transport", "high_railway_percent", "avg_time", "min_time",
				"avg_price", "min_price",
				"label_05060801", "label_05060701", "label_05060601", "label_02050601", "label_02050501", "label_02050401",
				"is_match_category", "train_consumer_prefer", "flight_consumer_prefer"
				, "bus_consumer_prefer"
		};

		VectorAssembler assembler = new VectorAssembler().setInputCols(features).setOutputCol("features");
		Map<String,Object> javaMap=new HashMap<>();
		javaMap.put("objective","binary:logistic");
		javaMap.put("eta",0.1);
		javaMap.put("max_depth",7);
		javaMap.put("min_child_weight",1);
		javaMap.put("alpha",1);
		javaMap.put("eval_metric","logloss");
		javaMap.put("num_round","20");
		javaMap.put("missing",-1);
		javaMap.put("num_workers",4);
		javaMap.put("num_early_stopping_rounds",20);
		javaMap.put("maximize_evaluation_metrics",false);
		javaMap.put("silent",1);
		javaMap.put("seed",2019L);
		CommonUtil.toScalaImmutableMap(javaMap);
		XGBoostClassifier xgBoostEstimator=new XGBoostClassifier( CommonUtil.<String,Object>toScalaImmutableMap(javaMap))
				.setFeaturesCol("features").setLabelCol("isclick").setProbabilityCol("probabilities");

		Pipeline pipeline = new Pipeline()
				.setStages(new PipelineStage[]{assembler,xgBoostEstimator});

		PipelineModel pipelineModel = pipeline.fit(trainData);

		Dataset<Row> predictResult=pipelineModel.transform(trainData);
		predictResult.show(false);

		XGBoostClassificationModel xgBoostClassificationModel=(XGBoostClassificationModel) (pipelineModel.stages()[1]);

		System.out.println(xgBoostClassificationModel.extractParamMap());
		//evaluate
		BinaryClassificationEvaluator evaluator=new BinaryClassificationEvaluator().setLabelCol("isclick").setRawPredictionCol("probabilities");
//				.setRawPredictionCol("probability");
		Double aucArea=evaluator.evaluate(predictResult);
		System.out.println("auc is :"+aucArea);

		pipelineModel.write().overwrite().save("/data/twms/traffichuixing/model/xgboost");
		savePMML(trainData.schema(),pipelineModel);

	}

	private static void savePMML(StructType shcema, PipelineModel pipelineModel) throws IOException, JAXBException {
		PMML pmml = new PMMLBuilder(shcema, pipelineModel).build();
		String targetFile = "/data/twms/traffichuixing/model/xgboost-pmml";
        PMMLUtil.marshal(pmml, new FileOutputStream(targetFile));
//		PMMLUtil.marshal(pmml, HDFSFileSystem.fileSystem.create(new Path(targetFile)));
	}

}

我这demo运行在本地的。

依赖自定义方法
xgboost目前只支持scala map，所以写了CommonUtil.toScalaImmutableMap

package com.demo.util;

import scala.Tuple2;

/**
 * Java-HashMap转Scala-Immutable.Map
 *
 * @author AllenBai
 */
public class CommonUtil {
	@SuppressWarnings("unchecked")
	public static <K, V> scala.collection.immutable.Map<K, V> toScalaImmutableMap(java.util.Map<K, V> javaMap) {
		final java.util.List<scala.Tuple2<K, V>> list = new java.util.ArrayList<>(javaMap.size());
		for (final java.util.Map.Entry<K, V> entry : javaMap.entrySet()) {
			list.add(scala.Tuple2.<K, V>apply(entry.getKey(), entry.getValue()));
		}
		final scala.collection.Seq<Tuple2<K, V>> seq = scala.collection.JavaConverters.asScalaBufferConverter(list).asScala().toSeq();
		return (scala.collection.immutable.Map<K, V>) scala.collection.immutable.Map$.MODULE$.apply(seq);
	}
}

注意点：
如果设置seed，必须设置为Long类型。
要不会报save model的错误；0.72没这种错误。参考文档

Exception in thread "main" java.lang.ClassCastException: java.lang.Integer cannot be cast to java.lang.Long at scala.runtime.BoxesRunTime.unboxToLong(BoxesRunTime.java:105)

---

项目的github地址xgboost-spark