Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[jvm-packages] For training data with group, empty RDD partition threw exception (#3749) (#3750)

Browse Source
This commit is contained in:
weitian
2018-10-09 09:03:22 -07:00
committed by Nan Zhu
parent ca33bf6476
commit 9504f411c1
3 changed files with 12 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostGeneralSuite.scala
View File

@@ -24,6 +24,7 @@ import ml.dmlc.xgboost4j.scala.DMatrix
import ml.dmlc.xgboost4j.scala.rabit.RabitTracker
import ml.dmlc.xgboost4j.scala.{XGBoost => SXGBoost, _}
import org.apache.hadoop.fs.{FileSystem, Path}
import org.apache.spark.TaskContext
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.sql._
import org.scalatest.FunSuite
@@ -256,8 +257,16 @@ class XGBoostGeneralSuite extends FunSuite with PerTest {
}
}
test("repartitionForTrainingGroup with group data which has empty partition") {
val trainingRDD = sc.parallelize(Ranking.train, 5).mapPartitions(it => {
// make one partition empty for testing
it.filter(_ => TaskContext.getPartitionId() != 3)
})
XGBoost.repartitionForTrainingGroup(trainingRDD, 4)
}
test("distributed training with group data") {
val trainingRDD = sc.parallelize(Ranking.train, 2)
val trainingRDD = sc.parallelize(Ranking.train, 5)
val (booster, metrics) = XGBoost.trainDistributed(
trainingRDD,
List("eta" -> "1", "max_depth" -> "6", "silent" -> "1",