Fix time field extraction after upstream change (#873)

Elasticsearch changed doc_values of date fields to return a
joda DateTime object. Thus, we need to call getMillis() to extract
the epoch millis value.

Original commit: elastic/x-pack-elasticsearch@b992882

Author: dimitris-athanasiou

elasticsearch/src/main/java/org/elasticsearch/xpack/ml/datafeed/extractor/scroll/ExtractedField.java

@@ -7,6 +7,7 @@
 
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.SearchHitField;
+import org.joda.time.base.BaseDateTime;
 
 import java.util.List;
 import java.util.Map;
@@ -36,6 +37,13 @@ public ExtractionMethod getExtractionMethod() {
 
     public abstract Object[] value(SearchHit hit);
 
+    public static ExtractedField newTimeField(String name, ExtractionMethod extractionMethod) {
+        if (extractionMethod == ExtractionMethod.SOURCE) {
+            throw new IllegalArgumentException("time field cannot be extracted from source");
+        }
+        return new TimeField(name, extractionMethod);
+    }
+
     public static ExtractedField newField(String name, ExtractionMethod extractionMethod) {
         switch (extractionMethod) {
             case DOC_VALUE:
@@ -65,6 +73,23 @@ public Object[] value(SearchHit hit) {
         }
     }
 
+    private static class TimeField extends FromFields {
+
+        TimeField(String name, ExtractionMethod extractionMethod) {
+            super(name, extractionMethod);
+        }
+
+        @Override
+        public Object[] value(SearchHit hit) {
+            Object[] value = super.value(hit);
+            if (value.length != 1) {
+                return value;
+            }
+            value[0] = ((BaseDateTime) value[0]).getMillis();
+            return value;
+        }
+    }
+
     private static class FromSource extends ExtractedField {
 
         private String[] namePath;

elasticsearch/src/main/java/org/elasticsearch/xpack/ml/datafeed/extractor/scroll/ExtractedFields.java

@@ -71,7 +71,7 @@ public Long timeFieldValue(SearchHit hit) {
     public static ExtractedFields build(Job job, DatafeedConfig datafeedConfig) {
         Set<String> scriptFields = datafeedConfig.getScriptFields().stream().map(sf -> sf.fieldName()).collect(Collectors.toSet());
         String timeField = job.getDataDescription().getTimeField();
-        ExtractedField timeExtractedField = ExtractedField.newField(timeField, scriptFields.contains(timeField) ?
+        ExtractedField timeExtractedField = ExtractedField.newTimeField(timeField, scriptFields.contains(timeField) ?
                 ExtractedField.ExtractionMethod.SCRIPT_FIELD : ExtractedField.ExtractionMethod.DOC_VALUE);
         List<String> remainingFields = job.allFields().stream().filter(f -> !f.equals(timeField)).collect(Collectors.toList());
         List<ExtractedField> allExtractedFields = new ArrayList<>(remainingFields.size());

elasticsearch/src/test/java/org/elasticsearch/xpack/ml/datafeed/extractor/scroll/ExtractedFieldTests.java

@@ -11,6 +11,7 @@
 import org.elasticsearch.search.internal.InternalSearchHit;
 import org.elasticsearch.search.internal.InternalSearchHitField;
 import org.elasticsearch.test.ESTestCase;
+import org.joda.time.DateTime;
 
 import java.util.Arrays;
 import java.util.HashMap;
@@ -96,6 +97,18 @@ public void testValueGivenEmptyHit() {
         assertThat(sourceField.value(hit), equalTo(new Object[0]));
     }
 
+    public void testNewTimeFieldGivenSource() {
+        expectThrows(IllegalArgumentException.class, () -> ExtractedField.newTimeField("time", ExtractedField.ExtractionMethod.SOURCE));
+    }
+
+    public void testValueGivenTimeField() {
+        SearchHit hit = new SearchHitBuilder(42).addField("time", new DateTime(123456789L)).build();
+
+        ExtractedField timeField = ExtractedField.newTimeField("time", ExtractedField.ExtractionMethod.DOC_VALUE);
+
+        assertThat(timeField.value(hit), equalTo(new Object[] { 123456789L }));
+    }
+
     static class SearchHitBuilder {
 
         private final InternalSearchHit hit;

elasticsearch/src/test/java/org/elasticsearch/xpack/ml/datafeed/extractor/scroll/ExtractedFieldsTests.java

@@ -7,6 +7,7 @@
 
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.test.ESTestCase;
+import org.joda.time.DateTime;
 
 import java.util.Arrays;
 import java.util.Collections;
@@ -15,7 +16,7 @@
 
 public class ExtractedFieldsTests extends ESTestCase {
 
-    private ExtractedField timeField = ExtractedField.newField("time", ExtractedField.ExtractionMethod.DOC_VALUE);
+    private ExtractedField timeField = ExtractedField.newTimeField("time", ExtractedField.ExtractionMethod.DOC_VALUE);
 
     public void testInvalidConstruction() {
         expectThrows(IllegalArgumentException.class, () -> new ExtractedFields(timeField, Collections.emptyList()));
@@ -47,7 +48,7 @@ public void testAllTypesOfFields() {
     }
 
     public void testTimeFieldValue() {
-        SearchHit hit = new ExtractedFieldTests.SearchHitBuilder(1).addField("time", 1000L).build();
+        SearchHit hit = new ExtractedFieldTests.SearchHitBuilder(1).addField("time", new DateTime(1000L)).build();
 
         ExtractedFields extractedFields = new ExtractedFields(timeField, Arrays.asList(timeField));