[ML] Apply source query on data frame analytics memory estimation (#49517)

dimitris-athanasiou · web-flow · commit 1865d874e77f · 2019-11-25T11:51:49.000+02:00
Closes #49454
diff --git a/x-pack/plugin/ml/qa/native-multi-node-tests/src/test/java/org/elasticsearch/xpack/ml/integration/ExplainDataFrameAnalyticsIT.java b/x-pack/plugin/ml/qa/native-multi-node-tests/src/test/java/org/elasticsearch/xpack/ml/integration/ExplainDataFrameAnalyticsIT.java
@@ -0,0 +1,64 @@
+/*
+ * 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.ml.integration;
+
+import org.elasticsearch.action.bulk.BulkRequestBuilder;
+import org.elasticsearch.action.bulk.BulkResponse;
+import org.elasticsearch.action.index.IndexRequest;
+import org.elasticsearch.action.support.WriteRequest;
+import org.elasticsearch.index.query.QueryBuilders;
+import org.elasticsearch.xpack.core.ml.action.ExplainDataFrameAnalyticsAction;
+import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsConfig;
+import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsSource;
+import org.elasticsearch.xpack.core.ml.dataframe.analyses.Classification;
+import org.elasticsearch.xpack.core.ml.utils.QueryProvider;
+
+import java.io.IOException;
+
+import static org.hamcrest.Matchers.lessThanOrEqualTo;
+
+public class ExplainDataFrameAnalyticsIT extends MlNativeDataFrameAnalyticsIntegTestCase {
+
+    public void testSourceQueryIsApplied() throws IOException {
+        // To test the source query is applied when we extract data,
+        // we set up a job where we have a query which excludes all but one document.
+        // We then assert the memory estimation is low enough.
+
+        String sourceIndex = "test-source-query-is-applied";
+
+        client().admin().indices().prepareCreate(sourceIndex)
+            .addMapping("_doc", "numeric_1", "type=double", "numeric_2", "type=float", "categorical", "type=keyword")
+            .get();
+
+        BulkRequestBuilder bulkRequestBuilder = client().prepareBulk();
+        bulkRequestBuilder.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
+
+        for (int i = 0; i < 30; i++) {
+            IndexRequest indexRequest = new IndexRequest(sourceIndex);
+
+            // We insert one odd value out of 5 for one feature
+            indexRequest.source("numeric_1", 1.0, "numeric_2", 2.0, "categorical", i == 0 ? "only-one" : "normal");
+            bulkRequestBuilder.add(indexRequest);
+        }
+        BulkResponse bulkResponse = bulkRequestBuilder.get();
+        if (bulkResponse.hasFailures()) {
+            fail("Failed to index data: " + bulkResponse.buildFailureMessage());
+        }
+
+        String id = "test_source_query_is_applied";
+
+        DataFrameAnalyticsConfig config = new DataFrameAnalyticsConfig.Builder()
+            .setId(id)
+            .setSource(new DataFrameAnalyticsSource(new String[] { sourceIndex },
+                QueryProvider.fromParsedQuery(QueryBuilders.termQuery("categorical", "only-one"))))
+            .setAnalysis(new Classification("categorical"))
+            .buildForExplain();
+
+        ExplainDataFrameAnalyticsAction.Response explainResponse = explainDataFrame(config);
+
+        assertThat(explainResponse.getMemoryEstimation().getExpectedMemoryWithoutDisk().getKb(), lessThanOrEqualTo(500L));
+    }
+}
diff --git a/x-pack/plugin/ml/qa/native-multi-node-tests/src/test/java/org/elasticsearch/xpack/ml/integration/MlNativeDataFrameAnalyticsIntegTestCase.java b/x-pack/plugin/ml/qa/native-multi-node-tests/src/test/java/org/elasticsearch/xpack/ml/integration/MlNativeDataFrameAnalyticsIntegTestCase.java
@@ -20,6 +20,7 @@
 import org.elasticsearch.search.sort.SortOrder;
 import org.elasticsearch.xpack.core.ml.action.DeleteDataFrameAnalyticsAction;
 import org.elasticsearch.xpack.core.ml.action.EvaluateDataFrameAction;
+import org.elasticsearch.xpack.core.ml.action.ExplainDataFrameAnalyticsAction;
 import org.elasticsearch.xpack.core.ml.action.GetDataFrameAnalyticsAction;
 import org.elasticsearch.xpack.core.ml.action.GetDataFrameAnalyticsStatsAction;
 import org.elasticsearch.xpack.core.ml.action.PutDataFrameAnalyticsAction;
@@ -146,6 +147,11 @@ protected GetDataFrameAnalyticsStatsAction.Response.Stats getAnalyticsStats(Stri
         return stats.get(0);
     }
 
+    protected ExplainDataFrameAnalyticsAction.Response explainDataFrame(DataFrameAnalyticsConfig config) {
+        PutDataFrameAnalyticsAction.Request request = new PutDataFrameAnalyticsAction.Request(config);
+        return client().execute(ExplainDataFrameAnalyticsAction.INSTANCE, request).actionGet();
+    }
+
     protected EvaluateDataFrameAction.Response evaluateDataFrame(String index, Evaluation evaluation) {
         EvaluateDataFrameAction.Request request =
             new EvaluateDataFrameAction.Request()
@@ -156,12 +162,12 @@ protected EvaluateDataFrameAction.Response evaluateDataFrame(String index, Evalu
 
     protected static DataFrameAnalyticsConfig buildAnalytics(String id, String sourceIndex, String destIndex,
                                                              @Nullable String resultsField, DataFrameAnalysis analysis) {
-        DataFrameAnalyticsConfig.Builder configBuilder = new DataFrameAnalyticsConfig.Builder();
-        configBuilder.setId(id);
-        configBuilder.setSource(new DataFrameAnalyticsSource(new String[] { sourceIndex }, null));
-        configBuilder.setDest(new DataFrameAnalyticsDest(destIndex, resultsField));
-        configBuilder.setAnalysis(analysis);
-        return configBuilder.build();
+        return new DataFrameAnalyticsConfig.Builder()
+            .setId(id)
+            .setSource(new DataFrameAnalyticsSource(new String[] { sourceIndex }, null))
+            .setDest(new DataFrameAnalyticsDest(destIndex, resultsField))
+            .setAnalysis(analysis)
+            .build();
     }
 
     protected void assertIsStopped(String id) {
diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/dataframe/extractor/DataFrameDataExtractorFactory.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/dataframe/extractor/DataFrameDataExtractorFactory.java
@@ -25,15 +25,18 @@ public class DataFrameDataExtractorFactory {
     private final Client client;
     private final String analyticsId;
     private final List<String> indices;
+    private final QueryBuilder sourceQuery;
     private final ExtractedFields extractedFields;
     private final Map<String, String> headers;
     private final boolean includeRowsWithMissingValues;
 
-    public DataFrameDataExtractorFactory(Client client, String analyticsId, List<String> indices, ExtractedFields extractedFields,
-                                          Map<String, String> headers, boolean includeRowsWithMissingValues) {
+    private DataFrameDataExtractorFactory(Client client, String analyticsId, List<String> indices, QueryBuilder sourceQuery,
+                                         ExtractedFields extractedFields, Map<String, String> headers,
+                                         boolean includeRowsWithMissingValues) {
         this.client = Objects.requireNonNull(client);
         this.analyticsId = Objects.requireNonNull(analyticsId);
         this.indices = Objects.requireNonNull(indices);
+        this.sourceQuery = Objects.requireNonNull(sourceQuery);
         this.extractedFields = Objects.requireNonNull(extractedFields);
         this.headers = headers;
         this.includeRowsWithMissingValues = includeRowsWithMissingValues;
@@ -54,7 +57,12 @@ public DataFrameDataExtractor newExtractor(boolean includeSource) {
     }
 
     private QueryBuilder createQuery() {
-        return includeRowsWithMissingValues ? QueryBuilders.matchAllQuery() : allExtractedFieldsExistQuery();
+        BoolQueryBuilder query = QueryBuilders.boolQuery();
+        query.filter(sourceQuery);
+        if (includeRowsWithMissingValues == false) {
+            query.filter(allExtractedFieldsExistQuery());
+        }
+        return query;
     }
 
     private QueryBuilder allExtractedFieldsExistQuery() {
@@ -77,8 +85,8 @@ private QueryBuilder allExtractedFieldsExistQuery() {
      */
     public static DataFrameDataExtractorFactory createForSourceIndices(Client client, String taskId, DataFrameAnalyticsConfig config,
                                                                        ExtractedFields extractedFields) {
-        return new DataFrameDataExtractorFactory(client, taskId, Arrays.asList(config.getSource().getIndex()), extractedFields,
-            config.getHeaders(), config.getAnalysis().supportsMissingValues());
+        return new DataFrameDataExtractorFactory(client, taskId, Arrays.asList(config.getSource().getIndex()),
+            config.getSource().getParsedQuery(), extractedFields, config.getHeaders(), config.getAnalysis().supportsMissingValues());
     }
 
     /**
@@ -100,8 +108,8 @@ public static void createForDestinationIndex(Client client,
             extractedFieldsDetector -> {
                 ExtractedFields extractedFields = extractedFieldsDetector.detect().v1();
                 DataFrameDataExtractorFactory extractorFactory = new DataFrameDataExtractorFactory(client, config.getId(),
-                    Collections.singletonList(config.getDest().getIndex()), extractedFields, config.getHeaders(),
-                    config.getAnalysis().supportsMissingValues());
+                    Collections.singletonList(config.getDest().getIndex()), config.getSource().getParsedQuery(), extractedFields,
+                    config.getHeaders(), config.getAnalysis().supportsMissingValues());
                 listener.onResponse(extractorFactory);
             },
             listener::onFailure
diff --git a/x-pack/qa/rolling-upgrade/src/test/resources/rest-api-spec/test/mixed_cluster/90_ml_data_frame_analytics_crud.yml b/x-pack/qa/rolling-upgrade/src/test/resources/rest-api-spec/test/mixed_cluster/90_ml_data_frame_analytics_crud.yml
@@ -7,7 +7,7 @@
   - match: { count: 1 }
   - match: { data_frame_analytics.0.id: "old_cluster_outlier_detection_job" }
   - match: { data_frame_analytics.0.source.index: ["bwc_ml_outlier_detection_job_source"] }
-  - match: { data_frame_analytics.0.source.query: {"term" : { "user" : "Kimchy" }} }
+  - match: { data_frame_analytics.0.source.query: {"term" : { "user.keyword" : "Kimchy" }} }
   - match: { data_frame_analytics.0.dest.index: "old_cluster_outlier_detection_job_results" }
   - match: { data_frame_analytics.0.analysis: {
     "outlier_detection":{
@@ -56,7 +56,7 @@
   - match: { count: 1 }
   - match: { data_frame_analytics.0.id: "old_cluster_regression_job" }
   - match: { data_frame_analytics.0.source.index: ["bwc_ml_regression_job_source"] }
-  - match: { data_frame_analytics.0.source.query: {"term": { "user": "Kimchy" }} }
+  - match: { data_frame_analytics.0.source.query: {"term": { "user.keyword": "Kimchy" }} }
   - match: { data_frame_analytics.0.dest.index: "old_cluster_regression_job_results" }
   - match: { data_frame_analytics.0.analysis.regression.dependent_variable: "foo" }
   - match: { data_frame_analytics.0.analysis.regression.training_percent: 100.0 }
@@ -101,7 +101,7 @@
           {
             "source": {
               "index": "bwc_ml_outlier_detection_job_source",
-              "query": {"term" : { "user" : "Kimchy" }}
+              "query": {"term" : { "user.keyword" : "Kimchy" }}
             },
             "dest": {
               "index": "mixed_cluster_outlier_detection_job_results"
@@ -116,7 +116,7 @@
   - match: { count: 1 }
   - match: { data_frame_analytics.0.id: "mixed_cluster_outlier_detection_job" }
   - match: { data_frame_analytics.0.source.index: ["bwc_ml_outlier_detection_job_source"] }
-  - match: { data_frame_analytics.0.source.query: {"term": { "user": "Kimchy" }} }
+  - match: { data_frame_analytics.0.source.query: {"term": { "user.keyword": "Kimchy" }} }
   - match: { data_frame_analytics.0.dest.index: "mixed_cluster_outlier_detection_job_results" }
   - match: { data_frame_analytics.0.analysis: {
     "outlier_detection":{
diff --git a/x-pack/qa/rolling-upgrade/src/test/resources/rest-api-spec/test/old_cluster/90_ml_data_frame_analytics_crud.yml b/x-pack/qa/rolling-upgrade/src/test/resources/rest-api-spec/test/old_cluster/90_ml_data_frame_analytics_crud.yml
@@ -5,7 +5,8 @@ setup:
         index: bwc_ml_outlier_detection_job_source
         body: >
           {
-            "numeric_field_1": 42.0
+            "numeric_field_1": 42.0,
+            "user": "Kimchy"
           }
 
   - do:
@@ -14,7 +15,8 @@ setup:
         body: >
           {
             "numeric_field_1": 1.0,
-            "foo": 10.0
+            "foo": 10.0,
+            "user": "Kimchy"
           }
 
   - do:
@@ -31,7 +33,7 @@ setup:
           {
             "source": {
               "index": "bwc_ml_outlier_detection_job_source",
-              "query": {"term" : { "user" : "Kimchy" }}
+              "query": {"term" : { "user.keyword" : "Kimchy" }}
             },
             "dest": {
               "index": "old_cluster_outlier_detection_job_results"
@@ -50,7 +52,7 @@ setup:
           {
             "source": {
               "index": "bwc_ml_regression_job_source",
-              "query": {"term" : { "user" : "Kimchy" }}
+              "query": {"term" : { "user.keyword" : "Kimchy" }}
             },
             "dest": {
               "index": "old_cluster_regression_job_results"
diff --git a/x-pack/qa/rolling-upgrade/src/test/resources/rest-api-spec/test/upgraded_cluster/90_ml_data_frame_analytics_crud.yml b/x-pack/qa/rolling-upgrade/src/test/resources/rest-api-spec/test/upgraded_cluster/90_ml_data_frame_analytics_crud.yml
@@ -7,7 +7,7 @@
   - match: { count: 1 }
   - match: { data_frame_analytics.0.id: "old_cluster_outlier_detection_job" }
   - match: { data_frame_analytics.0.source.index: ["bwc_ml_outlier_detection_job_source"] }
-  - match: { data_frame_analytics.0.source.query: {"term": { "user": "Kimchy" }} }
+  - match: { data_frame_analytics.0.source.query: {"term": { "user.keyword": "Kimchy" }} }
   - match: { data_frame_analytics.0.dest.index: "old_cluster_outlier_detection_job_results" }
   - match: { data_frame_analytics.0.analysis: {
     "outlier_detection":{
@@ -36,7 +36,7 @@
   - match: { count: 1 }
   - match: { data_frame_analytics.0.id: "old_cluster_regression_job" }
   - match: { data_frame_analytics.0.source.index: ["bwc_ml_regression_job_source"] }
-  - match: { data_frame_analytics.0.source.query: {"term": { "user": "Kimchy" }} }
+  - match: { data_frame_analytics.0.source.query: {"term": { "user.keyword": "Kimchy" }} }
   - match: { data_frame_analytics.0.dest.index: "old_cluster_regression_job_results" }
   - match: { data_frame_analytics.0.analysis.regression.dependent_variable: "foo" }
   - match: { data_frame_analytics.0.analysis.regression.training_percent: 100.0 }
@@ -62,7 +62,7 @@
   - match: { count: 1 }
   - match: { data_frame_analytics.0.id: "mixed_cluster_outlier_detection_job" }
   - match: { data_frame_analytics.0.source.index: ["bwc_ml_outlier_detection_job_source"] }
-  - match: { data_frame_analytics.0.source.query: {"term": { "user": "Kimchy" }} }
+  - match: { data_frame_analytics.0.source.query: {"term": { "user.keyword": "Kimchy" }} }
   - match: { data_frame_analytics.0.dest.index: "mixed_cluster_outlier_detection_job_results" }
   - match: { data_frame_analytics.0.analysis: {
     "outlier_detection":{

Original file line number	Diff line number	Diff line change
`@@ -5,7 +5,8 @@ setup:`
`5`	`5`	`index: bwc_ml_outlier_detection_job_source`
`6`	`6`	`body: >`
`7`	`7`	`{`
`8`		`- "numeric_field_1": 42.0`
	`8`	`+ "numeric_field_1": 42.0,`
	`9`	`+ "user": "Kimchy"`
`9`	`10`	`}`
`10`	`11`
`11`	`12`	`- do:`
`@@ -14,7 +15,8 @@ setup:`
`14`	`15`	`body: >`
`15`	`16`	`{`
`16`	`17`	`"numeric_field_1": 1.0,`
`17`		`- "foo": 10.0`
	`18`	`+ "foo": 10.0,`
	`19`	`+ "user": "Kimchy"`
`18`	`20`	`}`
`19`	`21`
`20`	`22`	`- do:`
`@@ -31,7 +33,7 @@ setup:`
`31`	`33`	`{`
`32`	`34`	`"source": {`
`33`	`35`	`"index": "bwc_ml_outlier_detection_job_source",`
`34`		`- "query": {"term" : { "user" : "Kimchy" }}`
	`36`	`+ "query": {"term" : { "user.keyword" : "Kimchy" }}`
`35`	`37`	`},`
`36`	`38`	`"dest": {`
`37`	`39`	`"index": "old_cluster_outlier_detection_job_results"`
`@@ -50,7 +52,7 @@ setup:`
`50`	`52`	`{`
`51`	`53`	`"source": {`
`52`	`54`	`"index": "bwc_ml_regression_job_source",`
`53`		`- "query": {"term" : { "user" : "Kimchy" }}`
	`55`	`+ "query": {"term" : { "user.keyword" : "Kimchy" }}`
`54`	`56`	`},`
`55`	`57`	`"dest": {`
`56`	`58`	`"index": "old_cluster_regression_job_results"`