[FEATURE][ML] Write control message to signify end of data (#36158)

Author: dimitris-athanasiou

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

@@ -192,11 +192,6 @@ public List<String> getFieldNames() {
         return context.extractedFields.getAllFields().stream().map(ExtractedField::getAlias).collect(Collectors.toList());
     }
 
-    public String[] getFieldNamesArray() {
-        List<String> fieldNames = getFieldNames();
-        return fieldNames.toArray(new String[fieldNames.size()]);
-    }
-
     public DataSummary collectDataSummary() {
         SearchRequestBuilder searchRequestBuilder = new SearchRequestBuilder(client, SearchAction.INSTANCE)
             .setIndices(context.indices)
@@ -210,9 +205,9 @@ public DataSummary collectDataSummary() {
     public static class DataSummary {
 
         public final long rows;
-        public final long cols;
+        public final int cols;
 
-        public DataSummary(long rows, long cols) {
+        public DataSummary(long rows, int cols) {
             this.rows = rows;
             this.cols = cols;
         }

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

@@ -0,0 +1,38 @@
+/*
+ * 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.xpack.ml.process.writer.AbstractControlMsgWriter;
+import org.elasticsearch.xpack.ml.process.writer.LengthEncodedWriter;
+
+import java.io.IOException;
+
+public class AnalyticsControlMessageWriter extends AbstractControlMsgWriter {
+
+    /**
+     * This must match the code defined in the api::CDataFrameAnalyzer C++ class.
+     * The constant there is referred as RUN_ANALYSIS_CONTROL_MESSAGE_FIELD_VALUE
+     * 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";
+
+    /**
+     * Construct the control message writer with a LengthEncodedWriter
+     *
+     * @param lengthEncodedWriter The writer
+     * @param numberOfFields      The number of fields the process expects in each record
+     */
+    public AnalyticsControlMessageWriter(LengthEncodedWriter lengthEncodedWriter, int numberOfFields) {
+        super(lengthEncodedWriter, numberOfFields);
+    }
+
+    public void writeEndOfData() throws IOException {
+        writeMessage(END_OF_DATA_MESSAGE_CODE);
+        fillCommandBuffer();
+        lengthEncodedWriter.flush();
+    }
+}

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

@@ -7,5 +7,14 @@
 
 import org.elasticsearch.xpack.ml.process.NativeProcess;
 
+import java.io.IOException;
+
 public interface AnalyticsProcess extends NativeProcess {
+
+    /**
+     * Writes a control message that informs the process
+     * all data has been sent
+     * @throws IOException If an error occurs writing to the process
+     */
+    void writeEndOfDataMessage() throws IOException;
 }

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

@@ -22,20 +22,24 @@ public class AnalyticsProcessConfig implements ToXContentObject {
     private static final String ANALYSIS = "analysis";
 
     private final long rows;
-    private final long cols;
+    private final int cols;
     private final ByteSizeValue memoryLimit;
     private final int threads;
     private final DataFrameAnalysis analysis;
 
 
-    public AnalyticsProcessConfig(long rows, long cols, ByteSizeValue memoryLimit, int threads, DataFrameAnalysis analysis) {
+    public AnalyticsProcessConfig(long rows, int cols, ByteSizeValue memoryLimit, int threads, DataFrameAnalysis analysis) {
         this.rows = rows;
         this.cols = cols;
         this.memoryLimit = Objects.requireNonNull(memoryLimit);
         this.threads = threads;
         this.analysis = Objects.requireNonNull(analysis);
     }
 
+    public int cols() {
+        return cols;
+    }
+
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();

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

@@ -7,6 +7,7 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.apache.logging.log4j.message.ParameterizedMessage;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.common.unit.ByteSizeUnit;
 import org.elasticsearch.common.unit.ByteSizeValue;
@@ -44,27 +45,16 @@ public void processData(String jobId, DataFrameDataExtractor dataExtractor) {
         threadPool.generic().execute(() -> {
             AnalyticsProcess process = createProcess(jobId, dataExtractor);
             try {
-                // Fake header
-                process.writeRecord(dataExtractor.getFieldNamesArray());
-
-                while (dataExtractor.hasNext()) {
-                    Optional<List<DataFrameDataExtractor.Row>> rows = dataExtractor.next();
-                    if (rows.isPresent()) {
-                        for (DataFrameDataExtractor.Row row : rows.get()) {
-                            String[] rowValues = row.getValues();
-                            if (rowValues != null) {
-                                process.writeRecord(rowValues);
-                            }
-                        }
-                    }
-                }
+                writeHeaderRecord(dataExtractor, process);
+                writeDataRows(dataExtractor, process);
+                process.writeEndOfDataMessage();
                 process.flushStream();
 
                 LOGGER.debug("[{}] Closing process", jobId);
                 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();
@@ -75,6 +65,37 @@ public void processData(String jobId, DataFrameDataExtractor dataExtractor) {
         });
     }
 
+    private void writeDataRows(DataFrameDataExtractor dataExtractor, AnalyticsProcess process) throws IOException {
+        // The extra field is the control field (should be an empty string)
+        String[] record = new String[dataExtractor.getFieldNames().size() + 1];
+        // The value of the control field should be an empty string for data frame rows
+        record[record.length - 1] = "";
+
+        while (dataExtractor.hasNext()) {
+            Optional<List<DataFrameDataExtractor.Row>> rows = dataExtractor.next();
+            if (rows.isPresent()) {
+                for (DataFrameDataExtractor.Row row : rows.get()) {
+                    String[] rowValues = row.getValues();
+                    if (rowValues != null) {
+                        System.arraycopy(rowValues, 0, record, 0, rowValues.length);
+                        process.writeRecord(record);
+                    }
+                }
+            }
+        }
+    }
+
+    private void writeHeaderRecord(DataFrameDataExtractor dataExtractor, AnalyticsProcess process) throws IOException {
+        List<String> fieldNames = dataExtractor.getFieldNames();
+        String[] headerRecord = new String[fieldNames.size() + 1];
+        for (int i = 0; i < fieldNames.size(); i++) {
+            headerRecord[i] = fieldNames.get(i);
+        }
+        // The field name of the control field is dot
+        headerRecord[headerRecord.length - 1] = ".";
+        process.writeRecord(headerRecord);
+    }
+
     private AnalyticsProcess createProcess(String jobId, DataFrameDataExtractor dataExtractor) {
         // 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);

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

@@ -7,6 +7,7 @@
 
 import org.elasticsearch.xpack.ml.process.AbstractNativeProcess;
 
+import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.nio.file.Path;
@@ -31,4 +32,9 @@ public String getName() {
     public void persistState() {
         // Nothing to persist
     }
+
+    @Override
+    public void writeEndOfDataMessage() throws IOException {
+        new AnalyticsControlMessageWriter(recordWriter(), numberOfFields()).writeEndOfData();
+    }
 }

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

@@ -45,10 +45,13 @@ public AnalyticsProcess createAnalyticsProcess(String jobId, AnalyticsProcessCon
         ProcessPipes processPipes = new ProcessPipes(env, NAMED_PIPE_HELPER, AnalyticsBuilder.ANALYTICS, jobId,
                 true, false, true, true, false, false);
 
+        // The extra 1 is the control field
+        int numberOfFields = analyticsProcessConfig.cols() + 1;
+
         createNativeProcess(jobId, analyticsProcessConfig, filesToDelete, processPipes);
 
         NativeAnalyticsProcess analyticsProcess = new NativeAnalyticsProcess(jobId, processPipes.getLogStream().get(),
-                processPipes.getProcessInStream().get(), processPipes.getProcessOutStream().get(), null, 0,
+                processPipes.getProcessInStream().get(), processPipes.getProcessOutStream().get(), null, numberOfFields,
                 filesToDelete, () -> {});
 
 

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/analytics/process/AnalyticsControlMessageWriterTests.java

@@ -0,0 +1,50 @@
+/*
+ * 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.test.ESTestCase;
+import org.elasticsearch.xpack.ml.process.writer.LengthEncodedWriter;
+import org.junit.Before;
+import org.mockito.InOrder;
+import org.mockito.Mockito;
+
+import java.io.IOException;
+import java.util.stream.IntStream;
+
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+public class AnalyticsControlMessageWriterTests extends ESTestCase {
+
+    private LengthEncodedWriter lengthEncodedWriter;
+
+    @Before
+    public void setUpMocks() {
+        lengthEncodedWriter = Mockito.mock(LengthEncodedWriter.class);
+    }
+
+    public void testWriteEndOfData() throws IOException {
+        AnalyticsControlMessageWriter writer = new AnalyticsControlMessageWriter(lengthEncodedWriter, 4);
+
+        writer.writeEndOfData();
+
+        InOrder inOrder = inOrder(lengthEncodedWriter);
+        inOrder.verify(lengthEncodedWriter).writeNumFields(4);
+        inOrder.verify(lengthEncodedWriter, times(3)).writeField("");
+        inOrder.verify(lengthEncodedWriter).writeField("r");
+
+        StringBuilder spaces = new StringBuilder();
+        IntStream.rangeClosed(1, 8192).forEach(i -> spaces.append(' '));
+        inOrder.verify(lengthEncodedWriter).writeNumFields(4);
+        inOrder.verify(lengthEncodedWriter, times(3)).writeField("");
+        inOrder.verify(lengthEncodedWriter).writeField(spaces.toString());
+
+        inOrder.verify(lengthEncodedWriter).flush();
+
+        verifyNoMoreInteractions(lengthEncodedWriter);
+    }
+}