fix examples

jvm-packages/xgboost4j-example/src/main/scala/ml/dmlc/xgboost4j/scala/example/flink/DistTrainWithFlink.scala

@@ -34,7 +34,7 @@ object DistTrainWithFlink {
     // number of iterations
     val round = 2
     // train the model
-    val model = XGBoost.train(paramMap, trainData, round)
+    val model = XGBoost.train(trainData, paramMap, round)
     val predTest = model.predict(testData.map{x => x.vector})
     model.saveModelAsHadoopFile("file:///path/to/xgboost.model")
   }

jvm-packages/xgboost4j-example/src/main/scala/ml/dmlc/xgboost4j/scala/example/spark/DistTrainWithSpark.scala

@@ -16,29 +16,34 @@
 
 package ml.dmlc.xgboost4j.scala.example.spark
 
-import ml.dmlc.xgboost4j.scala.spark.XGBoost
+import ml.dmlc.xgboost4j.scala.DMatrix
+import ml.dmlc.xgboost4j.scala.spark.{DataUtils, XGBoost}
 import org.apache.spark.SparkContext
 import org.apache.spark.mllib.util.MLUtils
 
 object DistTrainWithSpark {
   def main(args: Array[String]): Unit = {
-    if (args.length != 4) {
+    if (args.length != 5) {
       println(
-        "usage: program num_of_rounds num_workers training_path model_path")
+        "usage: program num_of_rounds num_workers training_path test_path model_path")
       sys.exit(1)
     }
     val sc = new SparkContext()
     val inputTrainPath = args(2)
-    val outputModelPath = args(3)
+    val inputTestPath = args(3)
+    val outputModelPath = args(4)
     // number of iterations
     val numRound = args(0).toInt
-    val trainRDD = MLUtils.loadLibSVMFile(sc, inputTrainPath).repartition(args(1).toInt)
+    import DataUtils._
+    val trainRDD = MLUtils.loadLibSVMFile(sc, inputTrainPath)
+    val testSet = MLUtils.loadLibSVMFile(sc, inputTestPath).collect().iterator
     // training parameters
     val paramMap = List(
       "eta" -> 0.1f,
       "max_depth" -> 2,
       "objective" -> "binary:logistic").toMap
-    val xgboostModel = XGBoost.train(trainRDD, paramMap, numRound)
+    val xgboostModel = XGBoost.train(trainRDD, paramMap, numRound, nWorkers = args(1).toInt)
+    xgboostModel.predict(new DMatrix(testSet))
     // save model to HDFS path
     xgboostModel.saveModelAsHadoopFile(outputModelPath)
   }

jvm-packages/xgboost4j-flink/src/main/scala/ml/dmlc/xgboost4j/scala/flink/XGBoost.scala

@@ -81,12 +81,13 @@ object XGBoost {
   /**
     * Train a xgboost model with link.
     *
-    * @param params The parameters to XGBoost.
     * @param dtrain The training data.
+    * @param params The parameters to XGBoost.
     * @param round Number of rounds to train.
     */
-  def train(params: Map[String, Any],
+  def train(
       dtrain: DataSet[LabeledVector],
+      params: Map[String, Any],
       round: Int): XGBoostModel = {
     val tracker = new RabitTracker(dtrain.getExecutionEnvironment.getParallelism)
     if (tracker.start()) {

jvm-packages/xgboost4j-flink/src/main/scala/ml/dmlc/xgboost4j/scala/flink/XGBoostModel.scala

@@ -37,6 +37,15 @@ class XGBoostModel (booster: Booster) extends Serializable {
       .create(new Path(modelPath)))
   }
 
+  /**
+   * predict with the given DMatrix
+   * @param testSet the local test set represented as DMatrix
+   * @return prediction result
+   */
+  def predict(testSet: DMatrix): Array[Array[Float]] = {
+    booster.predict(testSet, true, 0)
+  }
+
   /**
     * Predict given vector dataset.
     *
@@ -44,7 +53,7 @@ class XGBoostModel (booster: Booster) extends Serializable {
     * @return The prediction result.
     */
   def predict(data: DataSet[Vector]) : DataSet[Array[Float]] = {
-    val predictMap: Iterator[Vector] => TraversableOnce[Array[Float]] =
+    val predictMap: Iterator[Vector] => Traversable[Array[Float]] =
       (it: Iterator[Vector]) => {
         val mapper = (x: Vector) => {
           val (index, value) = x.toSeq.unzip

jvm-packages/xgboost4j/src/main/scala/ml/dmlc/xgboost4j/scala/DMatrix.scala

@@ -35,10 +35,10 @@ class DMatrix private[scala](private[scala] val jDMatrix: JDMatrix) {
     *  init DMatrix from Iterator of LabeledPoint
     *
     * @param dataIter An iterator of LabeledPoint
-    * @param cacheInfo  Cache path information, used for external memory setting, can be null.
+    * @param cacheInfo  Cache path information, used for external memory setting, null by default.
     * @throws XGBoostError native error
     */
-  def this(dataIter: Iterator[LabeledPoint], cacheInfo: String) {
+  def this(dataIter: Iterator[LabeledPoint], cacheInfo: String = null) {
     this(new JDMatrix(dataIter.asJava, cacheInfo))
   }