[FEATURE][ML] Parse results and join them in the data-frame copy index (#36382)

Author: dimitris-athanasiou

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

@@ -38,7 +38,6 @@
 import org.elasticsearch.xpack.core.ml.action.RunAnalyticsAction;
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
 import org.elasticsearch.xpack.ml.MachineLearning;
-import org.elasticsearch.xpack.ml.analytics.DataFrameDataExtractor;
 import org.elasticsearch.xpack.ml.analytics.DataFrameDataExtractorFactory;
 import org.elasticsearch.xpack.ml.analytics.DataFrameFields;
 import org.elasticsearch.xpack.ml.analytics.process.AnalyticsProcessManager;
@@ -178,8 +177,7 @@ private void runPipelineAnalytics(String index, ActionListener<AcknowledgedRespo
 
         ActionListener<DataFrameDataExtractorFactory> dataExtractorFactoryListener = ActionListener.wrap(
             dataExtractorFactory -> {
-                DataFrameDataExtractor dataExtractor = dataExtractorFactory.newExtractor();
-                analyticsProcessManager.processData(jobId, dataExtractor);
+                analyticsProcessManager.runJob(jobId, dataExtractorFactory);
                 listener.onResponse(new AcknowledgedResponse(true));
             },
             listener::onFailure

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/analytics/DataFrameDataExtractor.java

@@ -98,7 +98,7 @@ private SearchRequestBuilder buildSearchRequest() {
                 .setIndices(context.indices)
                 .setSize(context.scrollSize)
                 .setQuery(context.query)
-                .setFetchSource(false);
+                .setFetchSource(context.includeSource);
 
         for (ExtractedField docValueField : context.extractedFields.getDocValueFields()) {
             searchRequestBuilder.addDocValueField(docValueField.getName(), docValueField.getDocValueFormat());
@@ -149,7 +149,7 @@ private Row createRow(SearchHit hit) {
                 break;
             }
         }
-        return new Row(extractedValues);
+        return new Row(extractedValues, hit);
     }
 
     private List<Row> continueScroll() throws IOException {
@@ -196,10 +196,11 @@ public DataSummary collectDataSummary() {
         SearchRequestBuilder searchRequestBuilder = new SearchRequestBuilder(client, SearchAction.INSTANCE)
             .setIndices(context.indices)
             .setSize(0)
-            .setQuery(context.query);
+            .setQuery(context.query)
+            .setTrackTotalHits(true);
 
         SearchResponse searchResponse = executeSearchRequest(searchRequestBuilder);
-        return new DataSummary(searchResponse.getHits().getTotalHits(), context.extractedFields.getAllFields().size());
+        return new DataSummary(searchResponse.getHits().getTotalHits().value, context.extractedFields.getAllFields().size());
     }
 
     public static class DataSummary {
@@ -215,16 +216,27 @@ public DataSummary(long rows, int cols) {
 
     public static class Row {
 
+        private SearchHit hit;
+
         @Nullable
         private String[] values;
 
-        private Row(String[] values) {
+        private Row(String[] values, SearchHit hit) {
             this.values = values;
+            this.hit = hit;
         }
 
         @Nullable
         public String[] getValues() {
             return values;
         }
+
+        public SearchHit getHit() {
+            return hit;
+        }
+
+        public boolean shouldSkip() {
+            return values == null;
+        }
     }
 }

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/analytics/DataFrameDataExtractorContext.java

@@ -20,14 +20,16 @@ public class DataFrameDataExtractorContext {
     final QueryBuilder query;
     final int scrollSize;
     final Map<String, String> headers;
+    final boolean includeSource;
 
     DataFrameDataExtractorContext(String jobId, ExtractedFields extractedFields, List<String> indices, QueryBuilder query, int scrollSize,
-                                  Map<String, String> headers) {
+                                  Map<String, String> headers, boolean includeSource) {
         this.jobId = Objects.requireNonNull(jobId);
         this.extractedFields = Objects.requireNonNull(extractedFields);
         this.indices = indices.toArray(new String[indices.size()]);
         this.query = Objects.requireNonNull(query);
         this.scrollSize = scrollSize;
         this.headers = headers;
+        this.includeSource = includeSource;
     }
 }

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/analytics/DataFrameDataExtractorFactory.java

@@ -63,14 +63,16 @@ private DataFrameDataExtractorFactory(Client client, String index, ExtractedFiel
         this.extractedFields = Objects.requireNonNull(extractedFields);
     }
 
-    public DataFrameDataExtractor newExtractor() {
+    public DataFrameDataExtractor newExtractor(boolean includeSource) {
         DataFrameDataExtractorContext context = new DataFrameDataExtractorContext(
                 "ml-analytics-" + index,
                 extractedFields,
                 Arrays.asList(index),
                 QueryBuilders.matchAllQuery(),
                 1000,
-                Collections.emptyMap());
+                Collections.emptyMap(),
+                includeSource
+            );
         return new DataFrameDataExtractor(client, context);
     }
 

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/analytics/process/AnalyticsControlMessageWriter.java

@@ -18,7 +18,7 @@ public class AnalyticsControlMessageWriter extends AbstractControlMsgWriter {
      * but in the context of the java side it is more descriptive to call this the
      * end of data message.
      */
-    private static final String END_OF_DATA_MESSAGE_CODE = "r";
+    private static final String END_OF_DATA_MESSAGE_CODE = "$";
 
     /**
      * Construct the control message writer with a LengthEncodedWriter

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/analytics/process/AnalyticsProcess.java

@@ -8,6 +8,7 @@
 import org.elasticsearch.xpack.ml.process.NativeProcess;
 
 import java.io.IOException;
+import java.util.Iterator;
 
 public interface AnalyticsProcess extends NativeProcess {
 
@@ -17,4 +18,17 @@ public interface AnalyticsProcess extends NativeProcess {
      * @throws IOException If an error occurs writing to the process
      */
     void writeEndOfDataMessage() throws IOException;
+
+    /**
+     * @return stream of analytics results.
+     */
+    Iterator<AnalyticsResult> readAnalyticsResults();
+
+    /**
+     * Read anything left in the stream before
+     * closing the stream otherwise if the process
+     * tries to write more after the close it gets
+     * a SIGPIPE
+     */
+    void consumeAndCloseOutputStream();
 }

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/analytics/process/AnalyticsProcessManager.java

@@ -17,6 +17,7 @@
 import org.elasticsearch.xpack.ml.MachineLearning;
 import org.elasticsearch.xpack.ml.analytics.DataFrameAnalysis;
 import org.elasticsearch.xpack.ml.analytics.DataFrameDataExtractor;
+import org.elasticsearch.xpack.ml.analytics.DataFrameDataExtractorFactory;
 
 import java.io.IOException;
 import java.util.List;
@@ -41,28 +42,39 @@ public AnalyticsProcessManager(Client client, Environment environment, ThreadPoo
         this.processFactory = Objects.requireNonNull(analyticsProcessFactory);
     }
 
-    public void processData(String jobId, DataFrameDataExtractor dataExtractor) {
+    public void runJob(String jobId, DataFrameDataExtractorFactory dataExtractorFactory) {
         threadPool.generic().execute(() -> {
-            AnalyticsProcess process = createProcess(jobId, dataExtractor);
-            try {
-                writeHeaderRecord(dataExtractor, process);
-                writeDataRows(dataExtractor, process);
-                process.writeEndOfDataMessage();
-                process.flushStream();
+            DataFrameDataExtractor dataExtractor = dataExtractorFactory.newExtractor(false);
+            AnalyticsProcess process = createProcess(jobId, createProcessConfig(dataExtractor));
+            ExecutorService executorService = threadPool.executor(MachineLearning.AUTODETECT_THREAD_POOL_NAME);
+            AnalyticsResultProcessor resultProcessor = new AnalyticsResultProcessor(client, dataExtractorFactory.newExtractor(true));
+            executorService.execute(() -> resultProcessor.process(process));
+            executorService.execute(() -> processData(jobId, dataExtractor, process, resultProcessor));
+        });
+    }
+
+    private void processData(String jobId, DataFrameDataExtractor dataExtractor, AnalyticsProcess process,
+                             AnalyticsResultProcessor resultProcessor) {
+        try {
+            writeHeaderRecord(dataExtractor, process);
+            writeDataRows(dataExtractor, process);
+            process.writeEndOfDataMessage();
+            process.flushStream();
 
-                LOGGER.debug("[{}] Closing process", jobId);
+            LOGGER.info("[{}] Waiting for result processor to complete", jobId);
+            resultProcessor.awaitForCompletion();
+            LOGGER.info("[{}] Result processor has completed", jobId);
+        } catch (IOException e) {
+            LOGGER.error(new ParameterizedMessage("[{}] Error writing data to the process", jobId), e);
+        } finally {
+            LOGGER.info("[{}] Closing process", jobId);
+            try {
                 process.close();
                 LOGGER.info("[{}] Closed process", jobId);
             } catch (IOException e) {
-                LOGGER.error(new ParameterizedMessage("[{}] Error writing data to the process", jobId), e);
-            } finally {
-                try {
-                    process.close();
-                } catch (IOException e) {
-                    LOGGER.error("[{}] Error closing data frame analyzer process", jobId);
-                }
+                LOGGER.error("[{}] Error closing data frame analyzer process", jobId);
             }
-        });
+        }
     }
 
     private void writeDataRows(DataFrameDataExtractor dataExtractor, AnalyticsProcess process) throws IOException {
@@ -75,8 +87,8 @@ private void writeDataRows(DataFrameDataExtractor dataExtractor, AnalyticsProces
             Optional<List<DataFrameDataExtractor.Row>> rows = dataExtractor.next();
             if (rows.isPresent()) {
                 for (DataFrameDataExtractor.Row row : rows.get()) {
-                    String[] rowValues = row.getValues();
-                    if (rowValues != null) {
+                    if (row.shouldSkip() == false) {
+                        String[] rowValues = row.getValues();
                         System.arraycopy(rowValues, 0, record, 0, rowValues.length);
                         process.writeRecord(record);
                     }
@@ -96,10 +108,10 @@ private void writeHeaderRecord(DataFrameDataExtractor dataExtractor, AnalyticsPr
         process.writeRecord(headerRecord);
     }
 
-    private AnalyticsProcess createProcess(String jobId, DataFrameDataExtractor dataExtractor) {
+    private AnalyticsProcess createProcess(String jobId, AnalyticsProcessConfig analyticsProcessConfig) {
         // TODO We should rename the thread pool to reflect its more general use now, e.g. JOB_PROCESS_THREAD_POOL_NAME
         ExecutorService executorService = threadPool.executor(MachineLearning.AUTODETECT_THREAD_POOL_NAME);
-        AnalyticsProcess process = processFactory.createAnalyticsProcess(jobId, createProcessConfig(dataExtractor), executorService);
+        AnalyticsProcess process = processFactory.createAnalyticsProcess(jobId, analyticsProcessConfig, executorService);
         if (process.isProcessAlive() == false) {
             throw ExceptionsHelper.serverError("Failed to start analytics process");
         }

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/analytics/process/AnalyticsResult.java

@@ -0,0 +1,73 @@
+/*
+ * 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.analytics.process;
+
+import org.elasticsearch.common.ParseField;
+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.Map;
+import java.util.Objects;
+
+public class AnalyticsResult implements ToXContentObject {
+
+    public static final ParseField TYPE = new ParseField("analytics_result");
+    public static final ParseField ID_HASH = new ParseField("id_hash");
+    public static final ParseField RESULTS = new ParseField("results");
+
+    static final ConstructingObjectParser<AnalyticsResult, Void> PARSER = new ConstructingObjectParser<>(TYPE.getPreferredName(),
+            a -> new AnalyticsResult((String) a[0], (Map<String, Object>) a[1]));
+
+    static {
+        PARSER.declareString(ConstructingObjectParser.constructorArg(), ID_HASH);
+        PARSER.declareObject(ConstructingObjectParser.constructorArg(), (p, context) -> p.map(), RESULTS);
+    }
+
+    private final String idHash;
+    private final Map<String, Object> results;
+
+    public AnalyticsResult(String idHash, Map<String, Object> results) {
+        this.idHash = Objects.requireNonNull(idHash);
+        this.results = Objects.requireNonNull(results);
+    }
+
+    public String getIdHash() {
+        return idHash;
+    }
+
+    public Map<String, Object> getResults() {
+        return results;
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject();
+        builder.field(ID_HASH.getPreferredName(), idHash);
+        builder.field(RESULTS.getPreferredName(), results);
+        builder.endObject();
+        return builder;
+    }
+
+    @Override
+    public boolean equals(Object other) {
+        if (this == other) {
+            return true;
+        }
+        if (other == null || getClass() != other.getClass()) {
+            return false;
+        }
+
+        AnalyticsResult that = (AnalyticsResult) other;
+        return Objects.equals(idHash, that.idHash) && Objects.equals(results, that.results);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(idHash, results);
+    }
+}