Make ModelSizeStats a result type (elastic/elasticsearch#500)

* Make ModelSizeStats a result type

* Address review comments

Original commit: elastic/x-pack-elasticsearch@74de7b36d8

elasticsearch/src/main/java/org/elasticsearch/xpack/prelert/job/ModelSizeStats.java

@@ -15,6 +15,7 @@ import org.elasticsearch.common.xcontent.ConstructingObjectParser;
 import org.elasticsearch.common.xcontent.ObjectParser.ValueType;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser.Token;
+import org.elasticsearch.xpack.prelert.job.results.Result;
 import org.elasticsearch.xpack.prelert.utils.time.TimeUtils;
 
 import java.io.IOException;
@@ -26,10 +27,15 @@ import java.util.Objects;
  */
 public class ModelSizeStats extends ToXContentToBytes implements Writeable {
 
+    /**
+     * Result type
+     */
+    public static final String RESULT_TYPE_VALUE = "model_size_stats";
+    public static final ParseField RESULT_TYPE_FIELD = new ParseField(RESULT_TYPE_VALUE);
+
     /**
      * Field Names
      */
-    private static final ParseField MODEL_SIZE_STATS_FIELD = new ParseField("model_size_stats");
     public static final ParseField MODEL_BYTES_FIELD = new ParseField("model_bytes");
     public static final ParseField TOTAL_BY_FIELD_COUNT_FIELD = new ParseField("total_by_field_count");
     public static final ParseField TOTAL_OVER_FIELD_COUNT_FIELD = new ParseField("total_over_field_count");
@@ -40,10 +46,11 @@ public class ModelSizeStats extends ToXContentToBytes implements Writeable {
     public static final ParseField TIMESTAMP_FIELD = new ParseField("timestamp");
 
     public static final ConstructingObjectParser<Builder, ParseFieldMatcherSupplier> PARSER = new ConstructingObjectParser<>(
-            MODEL_SIZE_STATS_FIELD.getPreferredName(), a -> new Builder((String) a[0]));
+            RESULT_TYPE_FIELD.getPreferredName(), a -> new Builder((String) a[0]));
 
     static {
         PARSER.declareString(ConstructingObjectParser.constructorArg(), Job.ID);
+        PARSER.declareString((modelSizeStat, s) -> {}, Result.RESULT_TYPE);
         PARSER.declareLong(Builder::setModelBytes, MODEL_BYTES_FIELD);
         PARSER.declareLong(Builder::setBucketAllocationFailuresCount, BUCKET_ALLOCATION_FAILURES_COUNT_FIELD);
         PARSER.declareLong(Builder::setTotalByFieldCount, TOTAL_BY_FIELD_COUNT_FIELD);
@@ -70,11 +77,6 @@ public class ModelSizeStats extends ToXContentToBytes implements Writeable {
         PARSER.declareField(Builder::setMemoryStatus, p -> MemoryStatus.fromString(p.text()), MEMORY_STATUS_FIELD, ValueType.STRING);
     }
 
-    /**
-     * Elasticsearch type
-     */
-    public static final ParseField TYPE = new ParseField("model_size_stats");
-
     /**
      * The status of the memory monitored by the ResourceMonitor. OK is default,
      * SOFT_LIMIT means that the models have done some aggressive pruning to
@@ -183,6 +185,7 @@ public class ModelSizeStats extends ToXContentToBytes implements Writeable {
 
     public XContentBuilder doXContentBody(XContentBuilder builder) throws IOException {
         builder.field(Job.ID.getPreferredName(), jobId);
+        builder.field(Result.RESULT_TYPE.getPreferredName(), RESULT_TYPE_VALUE);
         builder.field(MODEL_BYTES_FIELD.getPreferredName(), modelBytes);
         builder.field(TOTAL_BY_FIELD_COUNT_FIELD.getPreferredName(), totalByFieldCount);
         builder.field(TOTAL_OVER_FIELD_COUNT_FIELD.getPreferredName(), totalOverFieldCount);
@@ -283,7 +286,7 @@ public class ModelSizeStats extends ToXContentToBytes implements Writeable {
 
         public Builder(String jobId) {
             this.jobId = jobId;
-            id = TYPE.getPreferredName();
+            id = RESULT_TYPE_FIELD.getPreferredName();
             memoryStatus = MemoryStatus.OK;
             logTime = new Date();
         }

elasticsearch/src/main/java/org/elasticsearch/xpack/prelert/job/ModelSnapshot.java

@@ -63,7 +63,7 @@ public class ModelSnapshot extends ToXContentToBytes implements Writeable {
         PARSER.declareLong(ModelSnapshot::setRestorePriority, RESTORE_PRIORITY);
         PARSER.declareString(ModelSnapshot::setSnapshotId, SNAPSHOT_ID);
         PARSER.declareInt(ModelSnapshot::setSnapshotDocCount, SNAPSHOT_DOC_COUNT);
-        PARSER.declareObject(ModelSnapshot::setModelSizeStats, ModelSizeStats.PARSER, ModelSizeStats.TYPE);
+        PARSER.declareObject(ModelSnapshot::setModelSizeStats, ModelSizeStats.PARSER, ModelSizeStats.RESULT_TYPE_FIELD);
         PARSER.declareField(ModelSnapshot::setLatestRecordTimeStamp, p -> {
             if (p.currentToken() == Token.VALUE_NUMBER) {
                 return new Date(p.longValue());
@@ -173,7 +173,7 @@ public class ModelSnapshot extends ToXContentToBytes implements Writeable {
         }
         builder.field(SNAPSHOT_DOC_COUNT.getPreferredName(), snapshotDocCount);
         if (modelSizeStats != null) {
-            builder.field(ModelSizeStats.TYPE.getPreferredName(), modelSizeStats);
+            builder.field(ModelSizeStats.RESULT_TYPE_FIELD.getPreferredName(), modelSizeStats);
         }
         if (latestRecordTimeStamp != null) {
             builder.field(LATEST_RECORD_TIME.getPreferredName(), latestRecordTimeStamp.getTime());

elasticsearch/src/main/java/org/elasticsearch/xpack/prelert/job/persistence/ElasticsearchBatchedModelSizeStatsIterator.java

@@ -1,44 +0,0 @@
-/*
- * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
- * or more contributor license agreements. Licensed under the Elastic License;
- * you may not use this file except in compliance with the Elastic License.
- */
-package org.elasticsearch.xpack.prelert.job.persistence;
-
-import org.elasticsearch.ElasticsearchParseException;
-import org.elasticsearch.client.Client;
-import org.elasticsearch.common.ParseFieldMatcher;
-import org.elasticsearch.common.bytes.BytesReference;
-import org.elasticsearch.common.xcontent.XContentFactory;
-import org.elasticsearch.common.xcontent.XContentParser;
-import org.elasticsearch.search.SearchHit;
-
-import java.io.IOException;
-
-import org.elasticsearch.xpack.prelert.job.ModelSizeStats;
-
-public class ElasticsearchBatchedModelSizeStatsIterator extends ElasticsearchBatchedDocumentsIterator<ModelSizeStats> {
-    public ElasticsearchBatchedModelSizeStatsIterator(Client client, String jobId, ParseFieldMatcher parserFieldMatcher) {
-        super(client, JobResultsPersister.getJobIndexName(jobId), parserFieldMatcher);
-    }
-
-    @Override
-    protected String getType() {
-        return ModelSizeStats.TYPE.getPreferredName();
-    }
-
-    @Override
-    protected ModelSizeStats map(SearchHit hit) {
-        BytesReference source = hit.getSourceRef();
-        XContentParser parser;
-        try {
-            parser = XContentFactory.xContent(source).createParser(source);
-        } catch (IOException e) {
-            throw new ElasticsearchParseException("failed to parser model size stats", e);
-        }
-
-        ModelSizeStats.Builder result = ModelSizeStats.PARSER.apply(parser, () -> parseFieldMatcher);
-        result.setId(hit.getId());
-        return result.build();
-    }
-}

elasticsearch/src/main/java/org/elasticsearch/xpack/prelert/job/persistence/ElasticsearchMappings.java

@@ -201,6 +201,7 @@ public class ElasticsearchMappings {
 
         addAnomalyRecordFieldsToMapping(builder);
         addInfluencerFieldsToMapping(builder);
+        addModelSizeStatsFieldsToMapping(builder);
 
         if (termFieldNames != null) {
             ElasticsearchDotNotationReverser reverser = new ElasticsearchDotNotationReverser();
@@ -508,7 +509,7 @@ public class ElasticsearchMappings {
      * document or by searching for all documents of this type.
      */
     public static XContentBuilder modelSnapshotMapping() throws IOException {
-        return jsonBuilder()
+        XContentBuilder builder = jsonBuilder()
                 .startObject()
                 .startObject(ModelSnapshot.TYPE.getPreferredName())
                 .startObject(ALL)
@@ -540,33 +541,15 @@ public class ElasticsearchMappings {
                 .startObject(ModelSnapshot.SNAPSHOT_DOC_COUNT.getPreferredName())
                 .field(TYPE, INTEGER)
                 .endObject()
-                .startObject(ModelSizeStats.TYPE.getPreferredName())
+                .startObject(ModelSizeStats.RESULT_TYPE_FIELD.getPreferredName())
                 .startObject(PROPERTIES)
                 .startObject(Job.ID.getPreferredName())
                 .field(TYPE, KEYWORD)
-                .endObject()
+                .endObject();
-                .startObject(ModelSizeStats.MODEL_BYTES_FIELD.getPreferredName())
+
-                .field(TYPE, LONG)
+        addModelSizeStatsFieldsToMapping(builder);
-                .endObject()
+
-                .startObject(ModelSizeStats.TOTAL_BY_FIELD_COUNT_FIELD.getPreferredName())
+        builder.startObject(ES_TIMESTAMP)
-                .field(TYPE, LONG)
-                .endObject()
-                .startObject(ModelSizeStats.TOTAL_OVER_FIELD_COUNT_FIELD.getPreferredName())
-                .field(TYPE, LONG)
-                .endObject()
-                .startObject(ModelSizeStats.TOTAL_PARTITION_FIELD_COUNT_FIELD.getPreferredName())
-                .field(TYPE, LONG)
-                .endObject()
-                .startObject(ModelSizeStats.BUCKET_ALLOCATION_FAILURES_COUNT_FIELD.getPreferredName())
-                .field(TYPE, LONG)
-                .endObject()
-                .startObject(ModelSizeStats.MEMORY_STATUS_FIELD.getPreferredName())
-                .field(TYPE, KEYWORD)
-                .endObject()
-                .startObject(ES_TIMESTAMP)
-                .field(TYPE, DATE)
-                .endObject()
-                .startObject(ModelSizeStats.LOG_TIME_FIELD.getPreferredName())
                 .field(TYPE, DATE)
                 .endObject()
                 .endObject()
@@ -593,50 +576,40 @@ public class ElasticsearchMappings {
                 .endObject()
                 .endObject()
                 .endObject();
+
+        return builder;
     }
 
     /**
-     * Create the Elasticsearch mapping for {@linkplain ModelSizeStats}.
+     * {@link ModelSizeStats} fields to be added under the 'properties' section of the mapping
+     * @param builder Add properties to this builder
+     * @return builder
+     * @throws IOException On write error
      */
-    public static XContentBuilder modelSizeStatsMapping() throws IOException {
+    private static XContentBuilder addModelSizeStatsFieldsToMapping(XContentBuilder builder) throws IOException {
-        return jsonBuilder()
+        builder.startObject(ModelSizeStats.MODEL_BYTES_FIELD.getPreferredName())
-                .startObject()
+                    .field(TYPE, LONG)
-                .startObject(ModelSizeStats.TYPE.getPreferredName())
-                .startObject(ALL)
-                .field(ENABLED, false)
-                // analyzer must be specified even though _all is disabled
-                // because all types in the same index must have the same
-                // analyzer for a given field
-                .field(ANALYZER, WHITESPACE)
-                .endObject()
-                .startObject(PROPERTIES)
-                .startObject(ModelSizeStats.MODEL_BYTES_FIELD.getPreferredName())
-                .field(TYPE, LONG)
                 .endObject()
                 .startObject(ModelSizeStats.TOTAL_BY_FIELD_COUNT_FIELD.getPreferredName())
-                .field(TYPE, LONG)
+                    .field(TYPE, LONG)
                 .endObject()
                 .startObject(ModelSizeStats.TOTAL_OVER_FIELD_COUNT_FIELD.getPreferredName())
-                .field(TYPE, LONG)
+                    .field(TYPE, LONG)
                 .endObject()
                 .startObject(ModelSizeStats.TOTAL_PARTITION_FIELD_COUNT_FIELD.getPreferredName())
-                .field(TYPE, LONG)
+                    .field(TYPE, LONG)
                 .endObject()
                 .startObject(ModelSizeStats.BUCKET_ALLOCATION_FAILURES_COUNT_FIELD.getPreferredName())
-                .field(TYPE, LONG)
+                    .field(TYPE, LONG)
                 .endObject()
                 .startObject(ModelSizeStats.MEMORY_STATUS_FIELD.getPreferredName())
-                .field(TYPE, KEYWORD)
+                    .field(TYPE, KEYWORD)
                 .endObject()
-                .startObject(ES_TIMESTAMP)
+                    .startObject(ModelSizeStats.LOG_TIME_FIELD.getPreferredName())
                 .field(TYPE, DATE)
-                .endObject()
-                .startObject(ModelSizeStats.LOG_TIME_FIELD.getPreferredName())
-                .field(TYPE, DATE)
-                .endObject()
-                .endObject()
-                .endObject()
                 .endObject();
+
+        return builder;
     }
 
     /**

elasticsearch/src/main/java/org/elasticsearch/xpack/prelert/job/persistence/JobProvider.java

@@ -187,7 +187,6 @@ public class JobProvider {
             XContentBuilder quantilesMapping = ElasticsearchMappings.quantilesMapping();
             XContentBuilder modelStateMapping = ElasticsearchMappings.modelStateMapping();
             XContentBuilder modelSnapshotMapping = ElasticsearchMappings.modelSnapshotMapping();
-            XContentBuilder modelSizeStatsMapping = ElasticsearchMappings.modelSizeStatsMapping();
             XContentBuilder modelDebugMapping = ElasticsearchMappings.modelDebugOutputMapping(termFields);
             XContentBuilder dataCountsMapping = ElasticsearchMappings.dataCountsMapping();
 
@@ -201,7 +200,6 @@ public class JobProvider {
             createIndexRequest.mapping(Quantiles.TYPE.getPreferredName(), quantilesMapping);
             createIndexRequest.mapping(ModelState.TYPE.getPreferredName(), modelStateMapping);
             createIndexRequest.mapping(ModelSnapshot.TYPE.getPreferredName(), modelSnapshotMapping);
-            createIndexRequest.mapping(ModelSizeStats.TYPE.getPreferredName(), modelSizeStatsMapping);
             createIndexRequest.mapping(ModelDebugOutput.TYPE.getPreferredName(), modelDebugMapping);
             createIndexRequest.mapping(DataCounts.TYPE.getPreferredName(), dataCountsMapping);
 
@@ -865,20 +863,6 @@ public class JobProvider {
         return new ElasticsearchBatchedModelDebugOutputIterator(client, jobId, parseFieldMatcher);
     }
 
-    /**
-     * Returns a {@link BatchedDocumentsIterator} that allows querying
-     * and iterating over a number of ModelSizeStats of the given job
-     *
-     * @param jobId the id of the job for which model snapshots are requested
-     * @return a model snapshot {@link BatchedDocumentsIterator}
-     */
-    public BatchedDocumentsIterator<ModelSizeStats> newBatchedModelSizeStatsIterator(String jobId) {
-        return new ElasticsearchBatchedModelSizeStatsIterator(client, jobId, parseFieldMatcher);
-    }
-
-    /**
-     * Get the persisted quantiles state for the job
-     */
     public Optional<Quantiles> getQuantiles(String jobId) {
         String indexName = JobResultsPersister.getJobIndexName(jobId);
         try {
@@ -983,7 +967,7 @@ public class JobProvider {
             Object timestamp = hit.getSource().remove(ElasticsearchMappings.ES_TIMESTAMP);
             hit.getSource().put(ModelSnapshot.TIMESTAMP.getPreferredName(), timestamp);
 
-            Object o = hit.getSource().get(ModelSizeStats.TYPE.getPreferredName());
+            Object o = hit.getSource().get(ModelSizeStats.RESULT_TYPE_FIELD.getPreferredName());
             if (o instanceof Map) {
                 @SuppressWarnings("unchecked")
                 Map<String, Object> map = (Map<String, Object>) o;
@@ -1091,11 +1075,11 @@ public class JobProvider {
     public Optional<ModelSizeStats> modelSizeStats(String jobId) {
         String indexName = JobResultsPersister.getJobIndexName(jobId);
         try {
-            LOGGER.trace("ES API CALL: get ID " + ModelSizeStats.TYPE +
+            LOGGER.trace("ES API CALL: get result type {} ID {} from index {}",
-                    " type " + ModelSizeStats.TYPE + " from index " + indexName);
+                    ModelSizeStats.RESULT_TYPE_VALUE, ModelSizeStats.RESULT_TYPE_FIELD, indexName);
 
             GetResponse modelSizeStatsResponse = client.prepareGet(
-                    indexName, ModelSizeStats.TYPE.getPreferredName(), ModelSizeStats.TYPE.getPreferredName()).get();
+                    indexName, Result.TYPE.getPreferredName(), ModelSizeStats.RESULT_TYPE_FIELD.getPreferredName()).get();
 
             if (!modelSizeStatsResponse.isExists()) {
                 String msg = "No memory usage details for job with id " + jobId;

elasticsearch/src/main/java/org/elasticsearch/xpack/prelert/job/persistence/JobResultsPersister.java

@@ -268,10 +268,10 @@ public class JobResultsPersister extends AbstractComponent {
     public void persistModelSizeStats(ModelSizeStats modelSizeStats) {
         String jobId = modelSizeStats.getJobId();
         logger.trace("[{}] Persisting model size stats, for size {}", jobId, modelSizeStats.getModelBytes());
-        Persistable persistable = new Persistable(modelSizeStats.getJobId(), modelSizeStats, ModelSizeStats.TYPE::getPreferredName,
+        Persistable persistable = new Persistable(modelSizeStats.getJobId(), modelSizeStats, Result.TYPE::getPreferredName,
-                ModelSizeStats.TYPE::getPreferredName, () -> toXContentBuilder(modelSizeStats));
+                ModelSizeStats.RESULT_TYPE_FIELD::getPreferredName, () -> toXContentBuilder(modelSizeStats));
         persistable.persist();
-        persistable = new Persistable(modelSizeStats.getJobId(), modelSizeStats, ModelSizeStats.TYPE::getPreferredName,
+        persistable = new Persistable(modelSizeStats.getJobId(), modelSizeStats, Result.TYPE::getPreferredName,
                 () -> null, () -> toXContentBuilder(modelSizeStats));
         persistable.persist();
         // Don't commit as we expect masses of these updates and they're only

elasticsearch/src/main/java/org/elasticsearch/xpack/prelert/job/results/AutodetectResult.java

@@ -39,7 +39,8 @@ public class AutodetectResult extends ToXContentToBytes implements Writeable {
         PARSER.declareObjectArray(ConstructingObjectParser.optionalConstructorArg(), Influencer.PARSER, Influencer.RESULTS_FIELD);
         PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(), Quantiles.PARSER, Quantiles.TYPE);
         PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(), ModelSnapshot.PARSER, ModelSnapshot.TYPE);
-        PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(), ModelSizeStats.PARSER, ModelSizeStats.TYPE);
+        PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(), ModelSizeStats.PARSER,
+                ModelSizeStats.RESULT_TYPE_FIELD);
         PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(), ModelDebugOutput.PARSER, ModelDebugOutput.TYPE);
         PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(), CategoryDefinition.PARSER, CategoryDefinition.TYPE);
         PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(), FlushAcknowledgement.PARSER, FlushAcknowledgement.TYPE);
@@ -154,7 +155,7 @@ public class AutodetectResult extends ToXContentToBytes implements Writeable {
         addNullableField(Influencer.RESULTS_FIELD, influencers, builder);
         addNullableField(Quantiles.TYPE, quantiles, builder);
         addNullableField(ModelSnapshot.TYPE, modelSnapshot, builder);
-        addNullableField(ModelSizeStats.TYPE, modelSizeStats, builder);
+        addNullableField(ModelSizeStats.RESULT_TYPE_FIELD, modelSizeStats, builder);
         addNullableField(ModelDebugOutput.TYPE, modelDebugOutput, builder);
         addNullableField(CategoryDefinition.TYPE, categoryDefinition, builder);
         addNullableField(FlushAcknowledgement.TYPE, flushAcknowledgement, builder);

elasticsearch/src/test/java/org/elasticsearch/xpack/prelert/job/persistence/ElasticsearchMappingsTests.java

@@ -93,8 +93,8 @@ public class ElasticsearchMappingsTests extends ESTestCase {
         overridden.add(ListDocument.TYPE.getPreferredName());
         overridden.add(ModelDebugOutput.TYPE.getPreferredName());
         overridden.add(ModelState.TYPE.getPreferredName());
+        overridden.add(ModelSizeStats.RESULT_TYPE_FIELD.getPreferredName());
         overridden.add(ModelSnapshot.TYPE.getPreferredName());
-        overridden.add(ModelSizeStats.TYPE.getPreferredName());
         overridden.add(Quantiles.TYPE.getPreferredName());
         overridden.add(Usage.TYPE);
 
@@ -147,11 +147,6 @@ public class ElasticsearchMappingsTests extends ESTestCase {
         parser = new JsonFactory().createParser(inputStream);
         parseJson(parser, expected);
 
-        builder = ElasticsearchMappings.modelSizeStatsMapping();
-        inputStream = new BufferedInputStream(new ByteArrayInputStream(builder.string().getBytes(StandardCharsets.UTF_8)));
-        parser = new JsonFactory().createParser(inputStream);
-        parseJson(parser, expected);
-
         builder = ElasticsearchMappings.modelSnapshotMapping();
         inputStream = new BufferedInputStream(new ByteArrayInputStream(builder.string().getBytes(StandardCharsets.UTF_8)));
         parser = new JsonFactory().createParser(inputStream);

elasticsearch/src/test/resources/rest-api-spec/test/jobs_get_stats.yaml

@@ -47,10 +47,11 @@ setup:
   - do:
       index:
         index:  prelertresults-job-stats-test
-        type:   model_size_stats
+        type:   result
         id:     model_size_stats
         body:   {
                   "job_id": "job-stats-test",
+                  "result_type": "model_size_stats",
                   "log_time": 1480896000000,
                   "timestamp": 1480896000000,
                   "model_bytes": 100,