From 4a2aba66c53af244652c770e26be98627d5ea99a Mon Sep 17 00:00:00 2001
From: Przemyslaw Witek <przemyslaw.witek@elastic.co>
Date: Mon, 5 Oct 2020 10:41:34 +0200
Subject: [PATCH] Make docCount field non-nullable

---
 .../classification/AbstractAucRoc.java           | 16 +++++++---------
 .../classification/AucRocResultTests.java        |  2 +-
 2 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/dataframe/evaluation/classification/AbstractAucRoc.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/dataframe/evaluation/classification/AbstractAucRoc.java
index 7e470cfbc276c..05f0323a509bb 100644
--- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/dataframe/evaluation/classification/AbstractAucRoc.java
+++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/dataframe/evaluation/classification/AbstractAucRoc.java
@@ -234,10 +234,10 @@ public static class Result implements EvaluationMetricResult {
         private static final String CURVE = "curve";
 
         private final double score;
-        private final Long docCount;
+        private final long docCount;
         private final List<AucRocPoint> curve;
 
-        public Result(double score, Long docCount, List<AucRocPoint> curve) {
+        public Result(double score, long docCount, List<AucRocPoint> curve) {
             this.score = score;
             this.docCount = docCount;
             this.curve = Objects.requireNonNull(curve);
@@ -245,7 +245,7 @@ public Result(double score, Long docCount, List<AucRocPoint> curve) {
 
         public Result(StreamInput in) throws IOException {
             this.score = in.readDouble();
-            this.docCount = in.readOptionalLong();
+            this.docCount = in.readLong();
             this.curve = in.readList(AucRocPoint::new);
         }
 
@@ -253,7 +253,7 @@ public double getScore() {
             return score;
         }
 
-        public Long getDocCount() {
+        public long getDocCount() {
             return docCount;
         }
 
@@ -274,7 +274,7 @@ public String getMetricName() {
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             out.writeDouble(score);
-            out.writeOptionalLong(docCount);
+            out.writeLong(docCount);
             out.writeList(curve);
         }
 
@@ -282,9 +282,7 @@ public void writeTo(StreamOutput out) throws IOException {
         public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
             builder.startObject();
             builder.field(SCORE, score);
-            if (docCount != null) {
-                builder.field(DOC_COUNT, docCount);
-            }
+            builder.field(DOC_COUNT, docCount);
             if (curve.isEmpty() == false) {
                 builder.field(CURVE, curve);
             }
@@ -298,7 +296,7 @@ public boolean equals(Object o) {
             if (o == null || getClass() != o.getClass()) return false;
             Result that = (Result) o;
             return score == that.score
-                && Objects.equals(docCount, that.docCount)
+                && docCount == that.docCount
                 && Objects.equals(curve, that.curve);
         }
 
diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/dataframe/evaluation/classification/AucRocResultTests.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/dataframe/evaluation/classification/AucRocResultTests.java
index 1a03e2e0c2c78..0c8a84aeba997 100644
--- a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/dataframe/evaluation/classification/AucRocResultTests.java
+++ b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/dataframe/evaluation/classification/AucRocResultTests.java
@@ -20,7 +20,7 @@ public class AucRocResultTests extends AbstractWireSerializingTestCase<Result> {
 
     public static Result createRandom() {
         double score = randomDoubleBetween(0.0, 1.0, true);
-        Long docCount = randomBoolean() ? randomLong() : null;
+        long docCount = randomLong();
         List<AucRocPoint> curve =
             Stream
                 .generate(() -> new AucRocPoint(randomDouble(), randomDouble(), randomDouble()))