allow empty partitions

2016-03-23 12:30:06 -04:00
parent f1114688a7
commit d8535313eb
2 changed files with 27 additions and 3 deletions
--- a/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostModel.scala
+++ b/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostModel.scala
@@ -32,8 +32,12 @@ class XGBoostModel(_booster: Booster)(implicit val sc: SparkContext) extends Ser
    import DataUtils._
    val broadcastBooster = testSet.sparkContext.broadcast(_booster)
    testSet.mapPartitions { testSamples =>
-      val dMatrix = new DMatrix(new JDMatrix(testSamples, null))
-      Iterator(broadcastBooster.value.predict(dMatrix))
+      if (testSamples.hasNext) {
+        val dMatrix = new DMatrix(new JDMatrix(testSamples, null))
+        Iterator(broadcastBooster.value.predict(dMatrix))
+      } else {
+        Iterator()
+      }
    }
  }