[jvm-packages] XGBoost Spark should deal with NaN when parsing evalua…

…tion output (#5546)
dmlc · Apr 20, 2020 · 397d8f0 · 397d8f0
1 parent b809f5d
commit 397d8f0
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/jvm-packages/xgboost4j/src/main/java/ml/dmlc/xgboost4j/java/Booster.java b/jvm-packages/xgboost4j/src/main/java/ml/dmlc/xgboost4j/java/Booster.java
@@ -242,7 +242,14 @@ public String evalSet(DMatrix[] evalMatrixs, String[] evalNames, int iter, float
     String stringFormat = evalSet(evalMatrixs, evalNames, iter);
     String[] metricPairs = stringFormat.split("\t");
     for (int i = 1; i < metricPairs.length; i++) {
-      metricsOut[i - 1] = Float.valueOf(metricPairs[i].split(":")[1]);
+      String value = metricPairs[i].split(":")[1];
+      if (value.equalsIgnoreCase("nan")) {
+        metricsOut[i - 1] = Float.NaN;
+      } else if (value.equalsIgnoreCase("-nan")) {
+        metricsOut[i - 1] = -Float.NaN;
+      } else {
+        metricsOut[i - 1] = Float.valueOf(value);
+      }
     }
     return stringFormat;
   }