[ML] Data frame analytics data counts

client/rest-high-level/src/main/java/org/elasticsearch/client/ml/dataframe/DataFrameAnalyticsStats.java

@@ -21,6 +21,7 @@
 
 import org.elasticsearch.client.ml.NodeAttributes;
 import org.elasticsearch.client.ml.dataframe.stats.AnalysisStats;
+import org.elasticsearch.client.ml.dataframe.stats.common.DataCounts;
 import org.elasticsearch.client.ml.dataframe.stats.common.MemoryUsage;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.ParseField;
@@ -47,6 +48,7 @@ public static DataFrameAnalyticsStats fromXContent(XContentParser parser) throws
     static final ParseField STATE = new ParseField("state");
     static final ParseField FAILURE_REASON = new ParseField("failure_reason");
     static final ParseField PROGRESS = new ParseField("progress");
+    static final ParseField DATA_COUNTS = new ParseField("data_counts");
     static final ParseField MEMORY_USAGE = new ParseField("memory_usage");
     static final ParseField ANALYSIS_STATS = new ParseField("analysis_stats");
     static final ParseField NODE = new ParseField("node");
@@ -60,10 +62,11 @@ public static DataFrameAnalyticsStats fromXContent(XContentParser parser) throws
                 (DataFrameAnalyticsState) args[1],
                 (String) args[2],
                 (List<PhaseProgress>) args[3],
-                (MemoryUsage) args[4],
-                (AnalysisStats) args[5],
-                (NodeAttributes) args[6],
-                (String) args[7]));
+                (DataCounts) args[4],
+                (MemoryUsage) args[5],
+                (AnalysisStats) args[6],
+                (NodeAttributes) args[7],
+                (String) args[8]));
 
     static {
         PARSER.declareString(constructorArg(), ID);
@@ -75,6 +78,7 @@ public static DataFrameAnalyticsStats fromXContent(XContentParser parser) throws
         }, STATE, ObjectParser.ValueType.STRING);
         PARSER.declareString(optionalConstructorArg(), FAILURE_REASON);
         PARSER.declareObjectArray(optionalConstructorArg(), PhaseProgress.PARSER, PROGRESS);
+        PARSER.declareObject(optionalConstructorArg(), DataCounts.PARSER, DATA_COUNTS);
         PARSER.declareObject(optionalConstructorArg(), MemoryUsage.PARSER, MEMORY_USAGE);
         PARSER.declareObject(optionalConstructorArg(), (p, c) -> parseAnalysisStats(p), ANALYSIS_STATS);
         PARSER.declareObject(optionalConstructorArg(), NodeAttributes.PARSER, NODE);
@@ -93,19 +97,21 @@ private static AnalysisStats parseAnalysisStats(XContentParser parser) throws IO
     private final DataFrameAnalyticsState state;
     private final String failureReason;
     private final List<PhaseProgress> progress;
+    private final DataCounts dataCounts;
     private final MemoryUsage memoryUsage;
     private final AnalysisStats analysisStats;
     private final NodeAttributes node;
     private final String assignmentExplanation;
 
     public DataFrameAnalyticsStats(String id, DataFrameAnalyticsState state, @Nullable String failureReason,
-                                   @Nullable List<PhaseProgress> progress, @Nullable MemoryUsage memoryUsage,
-                                   @Nullable AnalysisStats analysisStats, @Nullable NodeAttributes node,
+                                   @Nullable List<PhaseProgress> progress, @Nullable DataCounts dataCounts,
+                                   @Nullable MemoryUsage memoryUsage, @Nullable AnalysisStats analysisStats, @Nullable NodeAttributes node,
                                    @Nullable String assignmentExplanation) {
         this.id = id;
         this.state = state;
         this.failureReason = failureReason;
         this.progress = progress;
+        this.dataCounts = dataCounts;
         this.memoryUsage = memoryUsage;
         this.analysisStats = analysisStats;
         this.node = node;
@@ -128,6 +134,11 @@ public List<PhaseProgress> getProgress() {
         return progress;
     }
 
+    @Nullable
+    public DataCounts getDataCounts() {
+        return dataCounts;
+    }
+
     @Nullable
     public MemoryUsage getMemoryUsage() {
         return memoryUsage;
@@ -156,6 +167,7 @@ public boolean equals(Object o) {
             && Objects.equals(state, other.state)
             && Objects.equals(failureReason, other.failureReason)
             && Objects.equals(progress, other.progress)
+            && Objects.equals(dataCounts, other.dataCounts)
             && Objects.equals(memoryUsage, other.memoryUsage)
             && Objects.equals(analysisStats, other.analysisStats)
             && Objects.equals(node, other.node)
@@ -164,7 +176,7 @@ public boolean equals(Object o) {
 
     @Override
     public int hashCode() {
-        return Objects.hash(id, state, failureReason, progress, memoryUsage, analysisStats, node, assignmentExplanation);
+        return Objects.hash(id, state, failureReason, progress, dataCounts, memoryUsage, analysisStats, node, assignmentExplanation);
     }
 
     @Override
@@ -174,6 +186,7 @@ public String toString() {
             .add("state", state)
             .add("failureReason", failureReason)
             .add("progress", progress)
+            .add("dataCounts", dataCounts)
             .add("memoryUsage", memoryUsage)
             .add("analysisStats", analysisStats)
             .add("node", node)

client/rest-high-level/src/main/java/org/elasticsearch/client/ml/dataframe/stats/common/DataCounts.java

@@ -0,0 +1,119 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *         http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.client.ml.dataframe.stats.common;
+
+import org.elasticsearch.common.Nullable;
+import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.inject.internal.ToStringBuilder;
+import org.elasticsearch.common.xcontent.ConstructingObjectParser;
+import org.elasticsearch.common.xcontent.ToXContentObject;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+
+import java.io.IOException;
+import java.util.Objects;
+
+import static org.elasticsearch.common.xcontent.ConstructingObjectParser.optionalConstructorArg;
+
+public class DataCounts implements ToXContentObject {
+
+    public static final String TYPE_VALUE = "analytics_data_counts";
+
+    public static final ParseField TRAINING_DOCS_COUNT = new ParseField("training_docs_count");
+    public static final ParseField TEST_DOCS_COUNT = new ParseField("test_docs_count");
+    public static final ParseField SKIPPED_DOCS_COUNT = new ParseField("skipped_docs_count");
+
+    public static final ConstructingObjectParser<DataCounts, Void> PARSER = new ConstructingObjectParser<>(TYPE_VALUE, true,
+        a -> {
+            Long trainingDocsCount = (Long) a[0];
+            Long testDocsCount = (Long) a[1];
+            Long skippedDocsCount = (Long) a[2];
+            return new DataCounts(
+                getOrDefault(trainingDocsCount, 0L),
+                getOrDefault(testDocsCount, 0L),
+                getOrDefault(skippedDocsCount, 0L)
+            );
+        });
+
+    static {
+        PARSER.declareLong(optionalConstructorArg(), TRAINING_DOCS_COUNT);
+        PARSER.declareLong(optionalConstructorArg(), TEST_DOCS_COUNT);
+        PARSER.declareLong(optionalConstructorArg(), SKIPPED_DOCS_COUNT);
+    }
+
+    private final long trainingDocsCount;
+    private final long testDocsCount;
+    private final long skippedDocsCount;
+
+    public DataCounts(long trainingDocsCount, long testDocsCount, long skippedDocsCount) {
+        this.trainingDocsCount = trainingDocsCount;
+        this.testDocsCount = testDocsCount;
+        this.skippedDocsCount = skippedDocsCount;
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject();
+        builder.field(TRAINING_DOCS_COUNT.getPreferredName(), trainingDocsCount);
+        builder.field(TEST_DOCS_COUNT.getPreferredName(), testDocsCount);
+        builder.field(SKIPPED_DOCS_COUNT.getPreferredName(), skippedDocsCount);
+        builder.endObject();
+        return builder;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        DataCounts that = (DataCounts) o;
+        return trainingDocsCount == that.trainingDocsCount
+            && testDocsCount == that.testDocsCount
+            && skippedDocsCount == that.skippedDocsCount;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(trainingDocsCount, testDocsCount, skippedDocsCount);
+    }
+
+    @Override
+    public String toString() {
+        return new ToStringBuilder(getClass())
+            .add(TRAINING_DOCS_COUNT.getPreferredName(), trainingDocsCount)
+            .add(TEST_DOCS_COUNT.getPreferredName(), testDocsCount)
+            .add(SKIPPED_DOCS_COUNT.getPreferredName(), skippedDocsCount)
+            .toString();
+    }
+
+    public long getTrainingDocsCount() {
+        return trainingDocsCount;
+    }
+
+    public long getTestDocsCount() {
+        return testDocsCount;
+    }
+
+    public long getSkippedDocsCount() {
+        return skippedDocsCount;
+    }
+
+    private static <T> T getOrDefault(@Nullable T value, T defaultValue) {
+        return value != null ? value : defaultValue;
+    }
+}

client/rest-high-level/src/test/java/org/elasticsearch/client/ml/dataframe/DataFrameAnalyticsStatsTests.java

@@ -23,6 +23,7 @@
 import org.elasticsearch.client.ml.dataframe.stats.AnalysisStats;
 import org.elasticsearch.client.ml.dataframe.stats.AnalysisStatsNamedXContentProvider;
 import org.elasticsearch.client.ml.dataframe.stats.classification.ClassificationStatsTests;
+import org.elasticsearch.client.ml.dataframe.stats.common.DataCountsTests;
 import org.elasticsearch.client.ml.dataframe.stats.common.MemoryUsageTests;
 import org.elasticsearch.client.ml.dataframe.stats.outlierdetection.OutlierDetectionStatsTests;
 import org.elasticsearch.client.ml.dataframe.stats.regression.RegressionStatsTests;
@@ -68,6 +69,7 @@ public static DataFrameAnalyticsStats randomDataFrameAnalyticsStats() {
             randomFrom(DataFrameAnalyticsState.values()),
             randomBoolean() ? null : randomAlphaOfLength(10),
             randomBoolean() ? null : createRandomProgress(),
+            randomBoolean() ? null : DataCountsTests.createRandom(),
             randomBoolean() ? null : MemoryUsageTests.createRandom(),
             analysisStats,
             randomBoolean() ? null : NodeAttributesTests.createRandom(),
@@ -93,6 +95,9 @@ public static void toXContent(DataFrameAnalyticsStats stats, XContentBuilder bui
         if (stats.getProgress() != null) {
             builder.field(DataFrameAnalyticsStats.PROGRESS.getPreferredName(), stats.getProgress());
         }
+        if (stats.getDataCounts() != null) {
+            builder.field(DataFrameAnalyticsStats.DATA_COUNTS.getPreferredName(), stats.getDataCounts());
+        }
         if (stats.getMemoryUsage() != null) {
             builder.field(DataFrameAnalyticsStats.MEMORY_USAGE.getPreferredName(), stats.getMemoryUsage());
         }

client/rest-high-level/src/test/java/org/elasticsearch/client/ml/dataframe/stats/common/DataCountsTests.java

@@ -0,0 +1,51 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *         http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.client.ml.dataframe.stats.common;
+
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.test.AbstractXContentTestCase;
+
+import java.io.IOException;
+
+public class DataCountsTests extends AbstractXContentTestCase<DataCounts> {
+
+    @Override
+    protected DataCounts createTestInstance() {
+        return createRandom();
+    }
+
+    public static DataCounts createRandom() {
+        return new DataCounts(
+            randomNonNegativeLong(),
+            randomNonNegativeLong(),
+            randomNonNegativeLong()
+        );
+    }
+
+    @Override
+    protected DataCounts doParseInstance(XContentParser parser) throws IOException {
+        return DataCounts.PARSER.apply(parser, null);
+    }
+
+    @Override
+    protected boolean supportsUnknownFields() {
+        return true;
+    }
+}

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetDataFrameAnalyticsStatsAction.java

@@ -29,6 +29,7 @@
 import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsConfig;
 import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsState;
 import org.elasticsearch.xpack.core.ml.dataframe.stats.AnalysisStats;
+import org.elasticsearch.xpack.core.ml.dataframe.stats.common.DataCounts;
 import org.elasticsearch.xpack.core.ml.dataframe.stats.MemoryUsage;
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
 import org.elasticsearch.xpack.core.ml.utils.PhaseProgress;
@@ -165,6 +166,9 @@ public static class Stats implements ToXContentObject, Writeable {
              */
             private final List<PhaseProgress> progress;
 
+            @Nullable
+            private final DataCounts dataCounts;
+
             @Nullable
             private final MemoryUsage memoryUsage;
 
@@ -177,12 +181,13 @@ public static class Stats implements ToXContentObject, Writeable {
             private final String assignmentExplanation;
 
             public Stats(String id, DataFrameAnalyticsState state, @Nullable String failureReason, List<PhaseProgress> progress,
-                         @Nullable MemoryUsage memoryUsage, @Nullable AnalysisStats analysisStats, @Nullable DiscoveryNode node,
-                         @Nullable String assignmentExplanation) {
+                         @Nullable DataCounts dataCounts, @Nullable MemoryUsage memoryUsage, @Nullable AnalysisStats analysisStats,
+                         @Nullable DiscoveryNode node, @Nullable String assignmentExplanation) {
                 this.id = Objects.requireNonNull(id);
                 this.state = Objects.requireNonNull(state);
                 this.failureReason = failureReason;
                 this.progress = Objects.requireNonNull(progress);
+                this.dataCounts = dataCounts;
                 this.memoryUsage = memoryUsage;
                 this.analysisStats = analysisStats;
                 this.node = node;
@@ -198,6 +203,11 @@ public Stats(StreamInput in) throws IOException {
                 } else {
                     progress = in.readList(PhaseProgress::new);
                 }
+                if (in.getVersion().onOrAfter(Version.V_8_0_0)) {
+                    dataCounts = in.readOptionalWriteable(DataCounts::new);
+                } else {
+                    dataCounts = null;
+                }
                 if (in.getVersion().onOrAfter(Version.V_7_7_0)) {
                     memoryUsage = in.readOptionalWriteable(MemoryUsage::new);
                 } else {
@@ -261,6 +271,11 @@ public List<PhaseProgress> getProgress() {
                 return progress;
             }
 
+            @Nullable
+            public DataCounts getDataCounts() {
+                return dataCounts;
+            }
+
             @Nullable
             public MemoryUsage getMemoryUsage() {
                 return memoryUsage;
@@ -293,6 +308,9 @@ public XContentBuilder toUnwrappedXContent(XContentBuilder builder) throws IOExc
                 if (progress != null) {
                     builder.field("progress", progress);
                 }
+                if (dataCounts != null) {
+                    builder.field("data_counts", dataCounts);
+                }
                 if (memoryUsage != null) {
                     builder.field("memory_usage", memoryUsage);
                 }
@@ -331,6 +349,9 @@ public void writeTo(StreamOutput out) throws IOException {
                 } else {
                     out.writeList(progress);
                 }
+                if (out.getVersion().onOrAfter(Version.V_8_0_0)) {
+                    out.writeOptionalWriteable(dataCounts);
+                }
                 if (out.getVersion().onOrAfter(Version.V_7_7_0)) {
                     out.writeOptionalWriteable(memoryUsage);
                 }
@@ -369,7 +390,8 @@ private void writeProgressToLegacy(StreamOutput out) throws IOException {
 
             @Override
             public int hashCode() {
-                return Objects.hash(id, state, failureReason, progress, memoryUsage, analysisStats, node, assignmentExplanation);
+                return Objects.hash(id, state, failureReason, progress, dataCounts, memoryUsage, analysisStats, node,
+                    assignmentExplanation);
             }
 
             @Override
@@ -385,6 +407,7 @@ public boolean equals(Object obj) {
                         && Objects.equals(this.state, other.state)
                         && Objects.equals(this.failureReason, other.failureReason)
                         && Objects.equals(this.progress, other.progress)
+                        && Objects.equals(this.dataCounts, other.dataCounts)
                         && Objects.equals(this.memoryUsage, other.memoryUsage)
                         && Objects.equals(this.analysisStats, other.analysisStats)
                         && Objects.equals(this.node, other.node)