Uses MergingDigest instead of AVLDigest in percentiles agg

server/src/main/java/org/elasticsearch/search/aggregations/metrics/TDigestState.java

@@ -18,8 +18,9 @@
  */
 package org.elasticsearch.search.aggregations.metrics;
 
-import com.tdunning.math.stats.AVLTreeDigest;
 import com.tdunning.math.stats.Centroid;
+import com.tdunning.math.stats.MergingDigest;
+
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 
@@ -29,7 +30,7 @@
 /**
  * Extension of {@link com.tdunning.math.stats.TDigest} with custom serialization.
  */
-public class TDigestState extends AVLTreeDigest {
+public class TDigestState extends MergingDigest {
 
     private final double compression;
 
@@ -44,6 +45,8 @@ public double compression() {
     }
 
     public static void write(TDigestState state, StreamOutput out) throws IOException {
+        state.compress(); // This will flush any values in the buffer to the data structure. 
+                          // If there are no values in the buffer this is cheap
         out.writeDouble(state.compression);
         out.writeVInt(state.centroidCount());
         for (Centroid centroid : state.centroids()) {

server/src/test/java/org/elasticsearch/search/aggregations/metrics/InternalTDigestPercentilesRanksTests.java

@@ -42,7 +42,7 @@ protected InternalTDigestPercentileRanks createTestInstance(String name, List<Pi
         final TDigestState state = new TDigestState(100);
         Arrays.stream(values).forEach(state::add);
 
-        assertEquals(state.centroidCount(), values.length);
+        assertEquals(state.size(), values.length);
         return new InternalTDigestPercentileRanks(name, percents, state, keyed, format, aggregators, metadata);
     }
 

server/src/test/java/org/elasticsearch/search/aggregations/metrics/InternalTDigestPercentilesTests.java

@@ -43,7 +43,7 @@ protected InternalTDigestPercentiles createTestInstance(String name,
         final TDigestState state = new TDigestState(100);
         Arrays.stream(values).forEach(state::add);
 
-        assertEquals(state.centroidCount(), values.length);
+        assertEquals(state.size(), values.length);
         return new InternalTDigestPercentiles(name, percents, state, keyed, format, pipelineAggregators, metaData);
     }
 

server/src/test/java/org/elasticsearch/search/aggregations/metrics/TDigestPercentilesAggregatorTests.java

@@ -75,13 +75,13 @@ public void testSomeMatchesSortedNumericDocValues() throws IOException {
             iw.addDocument(singleton(new SortedNumericDocValuesField("number", 0)));
         }, tdigest -> {
             assertEquals(7L, tdigest.state.size());
-            assertEquals(7L, tdigest.state.centroidCount());
             assertEquals(4.5d, tdigest.percentile(75), 0.0d);
             assertEquals("4.5", tdigest.percentileAsString(75));
             assertEquals(2.0d, tdigest.percentile(50), 0.0d);
             assertEquals("2.0", tdigest.percentileAsString(50));
             assertEquals(1.0d, tdigest.percentile(22), 0.0d);
             assertEquals("1.0", tdigest.percentileAsString(22));
+            assertEquals(7L, tdigest.state.centroidCount());
         });
     }
 
@@ -96,7 +96,6 @@ public void testSomeMatchesNumericDocValues() throws IOException {
             iw.addDocument(singleton(new NumericDocValuesField("number", 0)));
         }, tdigest -> {
             assertEquals(tdigest.state.size(), 7L);
-            assertEquals(tdigest.state.centroidCount(), 7L);
             assertEquals(8.0d, tdigest.percentile(100), 0.0d);
             assertEquals("8.0", tdigest.percentileAsString(100));
             assertEquals(6.98d, tdigest.percentile(88), 0.0d);
@@ -107,6 +106,7 @@ public void testSomeMatchesNumericDocValues() throws IOException {
             assertEquals("1.0", tdigest.percentileAsString(25));
             assertEquals(0.0d, tdigest.percentile(1), 0.0d);
             assertEquals("0.0", tdigest.percentileAsString(1));
+            assertEquals(tdigest.state.centroidCount(), 7L);
         });
     }
 
@@ -123,10 +123,10 @@ public void testQueryFiltering() throws IOException {
 
         testCase(LongPoint.newRangeQuery("row", 1, 4), docs, tdigest -> {
             assertEquals(4L, tdigest.state.size());
-            assertEquals(4L, tdigest.state.centroidCount());
             assertEquals(2.0d, tdigest.percentile(100), 0.0d);
             assertEquals(1.0d, tdigest.percentile(50), 0.0d);
             assertEquals(0.5d, tdigest.percentile(25), 0.0d);
+            assertEquals(4L, tdigest.state.centroidCount());
         });
 
         testCase(LongPoint.newRangeQuery("row", 100, 110), docs, tdigest -> {