[ML] Add support for percentiles aggregation (elastic/x-pack-elasticsearch#733)

Original commit: elastic/x-pack-elasticsearch@7c8357ac04

plugin/src/main/java/org/elasticsearch/xpack/ml/datafeed/extractor/aggregation/AggregationToJsonProcessor.java

@@ -13,6 +13,8 @@ import org.elasticsearch.search.aggregations.Aggregations;
 import org.elasticsearch.search.aggregations.bucket.histogram.Histogram;
 import org.elasticsearch.search.aggregations.bucket.terms.Terms;
 import org.elasticsearch.search.aggregations.metrics.NumericMetricsAggregation;
+import org.elasticsearch.search.aggregations.metrics.percentiles.Percentile;
+import org.elasticsearch.search.aggregations.metrics.percentiles.Percentiles;
 import org.elasticsearch.xpack.ml.datafeed.DatafeedConfig;
 import org.joda.time.base.BaseDateTime;
 
@@ -20,6 +22,7 @@ import java.io.IOException;
 import java.io.OutputStream;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -79,7 +82,9 @@ class AggregationToJsonProcessor implements Releasable {
             List<String> addedKeys = new ArrayList<>();
             for (Aggregation nestedAgg : aggs) {
                 if (nestedAgg instanceof NumericMetricsAggregation.SingleValue) {
-                    addedKeys.add(processSingleValue(docCount, (NumericMetricsAggregation.SingleValue) nestedAgg));
+                    addedKeys.add(processSingleValue((NumericMetricsAggregation.SingleValue) nestedAgg));
+                } else if (nestedAgg instanceof Percentiles) {
+                    addedKeys.add(processPercentiles((Percentiles) nestedAgg));
                 } else {
                     throw new IllegalArgumentException("Unsupported aggregation type [" + nestedAgg.getName() + "]");
                 }
@@ -97,11 +102,20 @@ class AggregationToJsonProcessor implements Releasable {
         }
     }
 
-    private String processSingleValue(long docCount, NumericMetricsAggregation.SingleValue singleValue) throws IOException {
+    private String processSingleValue(NumericMetricsAggregation.SingleValue singleValue) throws IOException {
         keyValuePairs.put(singleValue.getName(), singleValue.value());
         return singleValue.getName();
     }
 
+    private String processPercentiles(Percentiles percentiles) throws IOException {
+        Iterator<Percentile> percentileIterator = percentiles.iterator();
+        keyValuePairs.put(percentiles.getName(), percentileIterator.next().getValue());
+        if (percentileIterator.hasNext()) {
+            throw new IllegalArgumentException("Multi-percentile aggregation [" + percentiles.getName() + "] is not supported");
+        }
+        return percentiles.getName();
+    }
+
     private void writeJsonObject(long docCount) throws IOException {
         if (docCount > 0) {
             jsonBuilder.startObject();

plugin/src/test/java/org/elasticsearch/xpack/ml/datafeed/extractor/aggregation/AggregationTestUtils.java

@@ -11,6 +11,8 @@ import org.elasticsearch.search.aggregations.bucket.histogram.Histogram;
 import org.elasticsearch.search.aggregations.bucket.terms.StringTerms;
 import org.elasticsearch.search.aggregations.bucket.terms.Terms;
 import org.elasticsearch.search.aggregations.metrics.NumericMetricsAggregation;
+import org.elasticsearch.search.aggregations.metrics.percentiles.Percentile;
+import org.elasticsearch.search.aggregations.metrics.percentiles.Percentiles;
 import org.joda.time.DateTime;
 
 import java.util.ArrayList;
@@ -83,6 +85,19 @@ public final class AggregationTestUtils {
         return termsAgg;
     }
 
+    static Percentiles createPercentiles(String name, double... values) {
+        Percentiles percentiles = mock(Percentiles.class);
+        when(percentiles.getName()).thenReturn(name);
+        List<Percentile> percentileList = new ArrayList<>();
+        for (double value : values) {
+            Percentile percentile = mock(Percentile.class);
+            when(percentile.getValue()).thenReturn(value);
+            percentileList.add(percentile);
+        }
+        when(percentiles.iterator()).thenReturn(percentileList.iterator());
+        return percentiles;
+    }
+
     static class Term {
         String key;
         long count;

plugin/src/test/java/org/elasticsearch/xpack/ml/datafeed/extractor/aggregation/AggregationToJsonProcessorTests.java

@@ -24,6 +24,7 @@ import static org.elasticsearch.xpack.ml.datafeed.extractor.aggregation.Aggregat
 import static org.elasticsearch.xpack.ml.datafeed.extractor.aggregation.AggregationTestUtils.createAggs;
 import static org.elasticsearch.xpack.ml.datafeed.extractor.aggregation.AggregationTestUtils.createDateHistogramBucket;
 import static org.elasticsearch.xpack.ml.datafeed.extractor.aggregation.AggregationTestUtils.createHistogramBucket;
+import static org.elasticsearch.xpack.ml.datafeed.extractor.aggregation.AggregationTestUtils.createPercentiles;
 import static org.elasticsearch.xpack.ml.datafeed.extractor.aggregation.AggregationTestUtils.createSingleValue;
 import static org.elasticsearch.xpack.ml.datafeed.extractor.aggregation.AggregationTestUtils.createTerms;
 import static org.hamcrest.Matchers.containsString;
@@ -217,6 +218,40 @@ public class AggregationToJsonProcessorTests extends ESTestCase {
         assertThat(json, equalTo("{\"time\":1000,\"doc_count\":3} {\"time\":2000,\"doc_count\":5}"));
     }
 
+    public void testProcessGivenSinglePercentilesPerHistogram() throws IOException {
+        List<Histogram.Bucket> histogramBuckets = Arrays.asList(
+                createHistogramBucket(1000L, 4, Arrays.asList(createPercentiles("my_field", 1.0))),
+                createHistogramBucket(2000L, 7, Arrays.asList(createPercentiles("my_field", 2.0))),
+                createHistogramBucket(3000L, 10, Arrays.asList(createPercentiles("my_field", 3.0))),
+                createHistogramBucket(4000L, 14, Arrays.asList(createPercentiles("my_field", 4.0)))
+        );
+        Histogram histogram = mock(Histogram.class);
+        when(histogram.getName()).thenReturn("time");
+        when(histogram.getBuckets()).thenReturn(histogramBuckets);
+
+        String json = aggToString(histogram);
+
+        assertThat(json, equalTo("{\"time\":1000,\"my_field\":1.0,\"doc_count\":4} " +
+                "{\"time\":2000,\"my_field\":2.0,\"doc_count\":7} " +
+                "{\"time\":3000,\"my_field\":3.0,\"doc_count\":10} " +
+                "{\"time\":4000,\"my_field\":4.0,\"doc_count\":14}"));
+    }
+
+    public void testProcessGivenMultiplePercentilesPerHistogram() throws IOException {
+        List<Histogram.Bucket> histogramBuckets = Arrays.asList(
+                createHistogramBucket(1000L, 4, Arrays.asList(createPercentiles("my_field", 1.0))),
+                createHistogramBucket(2000L, 7, Arrays.asList(createPercentiles("my_field", 2.0, 5.0))),
+                createHistogramBucket(3000L, 10, Arrays.asList(createPercentiles("my_field", 3.0))),
+                createHistogramBucket(4000L, 14, Arrays.asList(createPercentiles("my_field", 4.0)))
+        );
+        Histogram histogram = mock(Histogram.class);
+        when(histogram.getName()).thenReturn("time");
+        when(histogram.getBuckets()).thenReturn(histogramBuckets);
+
+        IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> aggToString(histogram));
+        assertThat(e.getMessage(), containsString("Multi-percentile aggregation [my_field] is not supported"));
+    }
+
     private String aggToString(Aggregation aggregation) throws IOException {
         return aggToString(aggregation, true);
     }