Add doc_count to ParsedMatrixStats (#24952)

This commit adds support in ParsedMatrixStats for parsing the doc_count
field.

Related to #24776

Author: tlrx

modules/aggs-matrix-stats/src/main/java/org/elasticsearch/search/aggregations/matrix/stats/ParsedMatrixStats.java

@@ -41,14 +41,20 @@ public class ParsedMatrixStats extends ParsedAggregation implements MatrixStats
     private final Map<String, Map<String, Double>> covariances = new HashMap<>();
     private final Map<String, Map<String, Double>> correlations = new HashMap<>();
 
+    private long docCount;
+
     @Override
     public String getType() {
         return MatrixStatsAggregationBuilder.NAME;
     }
 
+    private void setDocCount(long docCount) {
+        this.docCount = docCount;
+    }
+
     @Override
     public long getDocCount() {
-        throw new UnsupportedOperationException();
+        return docCount;
     }
 
     @Override
@@ -97,6 +103,7 @@ public double getCorrelation(String fieldX, String fieldY) {
 
     @Override
     protected XContentBuilder doXContentBody(XContentBuilder builder, Params params) throws IOException {
+        builder.field(CommonFields.DOC_COUNT.getPreferredName(), getDocCount());
         if (counts != null && counts.isEmpty() == false) {
             builder.startArray(InternalMatrixStats.Fields.FIELDS);
             for (String fieldName : counts.keySet()) {
@@ -148,6 +155,7 @@ private static <T> T checkedGet(final Map<String, T> values, final String fieldN
             new ObjectParser<>(ParsedMatrixStats.class.getSimpleName(), true, ParsedMatrixStats::new);
     static {
         declareAggregationFields(PARSER);
+        PARSER.declareLong(ParsedMatrixStats::setDocCount, CommonFields.DOC_COUNT);
         PARSER.declareObjectArray((matrixStats, results) -> {
             for (ParsedMatrixStatsResult result : results) {
                 final String fieldName = result.name;

modules/aggs-matrix-stats/src/test/java/org/elasticsearch/search/aggregations/matrix/stats/InternalMatrixStatsTests.java

@@ -133,8 +133,7 @@ protected void assertFromXContent(InternalMatrixStats expected, ParsedAggregatio
         assertTrue(parsedAggregation instanceof ParsedMatrixStats);
         ParsedMatrixStats actual = (ParsedMatrixStats) parsedAggregation;
 
-        //norelease add parsing logic for doc count and enable this test once elastic/elasticsearch#24776 is merged
-        //assertEquals(expected.getDocCount(), actual.getDocCount());
+        assertEquals(expected.getDocCount(), actual.getDocCount());
 
         for (String field : fields) {
             assertEquals(expected.getFieldCount(field), actual.getFieldCount(field));