Add serializable ExponentialHistogramState (#136075)

Author: JonasKunz

build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/InternalDistributionModuleCheckTaskProvider.java

@@ -53,6 +53,7 @@ public class InternalDistributionModuleCheckTaskProvider {
         "org.elasticsearch.base",
         "org.elasticsearch.cli",
         "org.elasticsearch.entitlement",
+        "org.elasticsearch.exponentialhistogram",
         "org.elasticsearch.geo",
         "org.elasticsearch.grok",
         "org.elasticsearch.logging",

libs/exponential-histogram/build.gradle

@@ -17,7 +17,9 @@ dependencies {
   api project(':libs:x-content')
   api "org.apache.lucene:lucene-core:${versions.lucene}"
 
-  testImplementation(project(":test:framework"))
+  testImplementation(project(":test:framework")) {
+    exclude group: 'org.elasticsearch', module: 'exponential-histogram'
+  }
   testImplementation('ch.obermuhlner:big-math:2.3.2')
   testImplementation('org.apache.commons:commons-math3:3.6.1')
   testImplementation project(':libs:x-content:impl')

libs/exponential-histogram/src/main/java/module-info.java

@@ -0,0 +1,26 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+module org.elasticsearch.exponentialhistogram {
+    requires org.elasticsearch.xcontent;
+    requires org.apache.lucene.core;
+    requires org.elasticsearch.base;
+
+    exports org.elasticsearch.exponentialhistogram;
+}

libs/exponential-histogram/src/main/java/org/elasticsearch/exponentialhistogram/ExponentialHistogram.java

@@ -270,7 +270,7 @@ static ReleasableExponentialHistogram create(int maxBucketCount, ExponentialHist
     static ReleasableExponentialHistogram merge(
         int maxBucketCount,
         ExponentialHistogramCircuitBreaker breaker,
-        Iterator<ExponentialHistogram> histograms
+        Iterator<? extends ExponentialHistogram> histograms
     ) {
         try (ExponentialHistogramMerger merger = ExponentialHistogramMerger.create(maxBucketCount, breaker)) {
             while (histograms.hasNext()) {

libs/exponential-histogram/src/main/java/org/elasticsearch/exponentialhistogram/ExponentialHistogramMerger.java

@@ -89,7 +89,11 @@ private ExponentialHistogramMerger(int bucketLimit, int maxScale, ExponentialHis
         downscaleStats = new DownscaleStats();
     }
 
-    static ExponentialHistogramMerger createForTesting(int bucketLimit, int maxScale, ExponentialHistogramCircuitBreaker circuitBreaker) {
+    public static ExponentialHistogramMerger createWithMaxScale(
+        int bucketLimit,
+        int maxScale,
+        ExponentialHistogramCircuitBreaker circuitBreaker
+    ) {
         circuitBreaker.adjustBreaker(BASE_SIZE);
         return new ExponentialHistogramMerger(bucketLimit, maxScale, circuitBreaker);
     }
@@ -137,6 +141,19 @@ public ReleasableExponentialHistogram getAndClear() {
         return retVal;
     }
 
+    /**
+     * Gets the current merged histogram without clearing this merger.
+     * Note that the ownership of the returned histogram remains with this merger,
+     * so the caller must not close it.
+     * The returned histogram is only valid until the next call to {@link #add(ExponentialHistogram)}, or until the merger is closed.
+     *
+     * @return the current merged histogram
+     */
+    public ExponentialHistogram get() {
+        assert closed == false : "ExponentialHistogramMerger already closed";
+        return (result == null) ? ExponentialHistogram.empty() : result;
+    }
+
     // This algorithm is very efficient if B has roughly as many buckets as A.
     // However, if B is much smaller we still have to iterate over all buckets of A.
     // This can be optimized by buffering the buckets of small histograms and only merging them when we have enough buckets.

libs/exponential-histogram/src/main/java/org/elasticsearch/exponentialhistogram/ZeroBucket.java

@@ -60,6 +60,11 @@ public final class ZeroBucket {
      */
     private double realThreshold;
 
+    /**
+     * True if the zero threshold was created from an index/scale pair, false if it was created from a real-valued threshold.
+     */
+    private final boolean isIndexBased;
+
     private final long count;
     // A singleton for an empty zero bucket with the smallest possible threshold.
     private static final ZeroBucket MINIMAL_EMPTY = new ZeroBucket(MIN_INDEX, MIN_SCALE, 0);
@@ -70,6 +75,7 @@ private ZeroBucket(double zeroThreshold, long count) {
         this.scale = MAX_SCALE;
         this.realThreshold = zeroThreshold;
         this.count = count;
+        this.isIndexBased = false;
     }
 
     private ZeroBucket(long index, int scale, long count) {
@@ -79,12 +85,14 @@ private ZeroBucket(long index, int scale, long count) {
         this.scale = scale;
         this.realThreshold = Double.NaN; // compute lazily when needed
         this.count = count;
+        this.isIndexBased = true;
     }
 
     private ZeroBucket(ZeroBucket toCopy, long newCount) {
         this.realThreshold = toCopy.realThreshold;
         this.index = toCopy.index;
         this.scale = toCopy.scale;
+        this.isIndexBased = toCopy.isIndexBased;
         this.count = newCount;
     }
 
@@ -163,6 +171,13 @@ public long count() {
         return count;
     }
 
+    /**
+     * @return True if the zero threshold was created from an index/scale pair, false if it was created from a real-valued threshold.
+     */
+    public boolean isIndexBased() {
+        return isIndexBased;
+    }
+
     /**
      * Merges this zero bucket with another one.
      * <ul>

libs/exponential-histogram/src/test/java/org/elasticsearch/exponentialhistogram/ExponentialHistogramMergerTests.java

@@ -43,7 +43,6 @@
 public class ExponentialHistogramMergerTests extends ExponentialHistogramTestCase {
 
     public void testZeroThresholdCollapsesOverlappingBuckets() {
-
         ExponentialHistogram first = createAutoReleasedHistogram(b -> b.zeroBucket(ZeroBucket.create(2.0001, 10)));
 
         ExponentialHistogram second = createAutoReleasedHistogram(
@@ -60,6 +59,7 @@ public void testZeroThresholdCollapsesOverlappingBuckets() {
 
         assertThat(mergeResult.zeroBucket().zeroThreshold(), equalTo(4.0));
         assertThat(mergeResult.zeroBucket().count(), equalTo(14L));
+        assertThat(mergeResult.zeroBucket().isIndexBased(), equalTo(true));
 
         // only the (4, 8] bucket should be left
         assertThat(mergeResult.scale(), equalTo(0));
@@ -82,6 +82,7 @@ public void testZeroThresholdCollapsesOverlappingBuckets() {
         mergeResult = mergeWithMinimumScale(100, 0, mergeResult, third);
         assertThat(mergeResult.zeroBucket().zeroThreshold(), closeTo(45.0, 0.000001));
         assertThat(mergeResult.zeroBucket().count(), equalTo(1L + 14L + 42L + 7L));
+        assertThat(mergeResult.zeroBucket().isIndexBased(), equalTo(false));
         assertThat(mergeResult.positiveBuckets().iterator().hasNext(), equalTo(false));
         assertThat(mergeResult.negativeBuckets().iterator().hasNext(), equalTo(false));
     }
@@ -241,7 +242,7 @@ private void assertBucketsEqual(ExponentialHistogram.Buckets bucketsA, Exponenti
     }
 
     private ExponentialHistogram mergeWithMinimumScale(int bucketCount, int scale, ExponentialHistogram... histograms) {
-        try (ExponentialHistogramMerger merger = ExponentialHistogramMerger.createForTesting(bucketCount, scale, breaker())) {
+        try (ExponentialHistogramMerger merger = ExponentialHistogramMerger.createWithMaxScale(bucketCount, scale, breaker())) {
             Arrays.stream(histograms).forEach(merger::add);
             ReleasableExponentialHistogram result = merger.getAndClear();
             autoReleaseOnTestEnd(result);

libs/exponential-histogram/src/test/java/org/elasticsearch/exponentialhistogram/ZeroBucketTests.java

@@ -33,6 +33,7 @@ public void testMinimalBucketHasZeroThreshold() {
     public void testExactThresholdPreserved() {
         ZeroBucket bucket = ZeroBucket.create(3.0, 10);
         assertThat(bucket.zeroThreshold(), equalTo(3.0));
+        assertThat(bucket.isIndexBased(), equalTo(false));
     }
 
     public void testMergingPreservesExactThreshold() {
@@ -41,6 +42,7 @@ public void testMergingPreservesExactThreshold() {
         ZeroBucket merged = bucketA.merge(bucketB);
         assertThat(merged.zeroThreshold(), equalTo(3.5));
         assertThat(merged.count(), equalTo(30L));
+        assertThat(merged.isIndexBased(), equalTo(false));
     }
 
     public void testBucketCollapsingPreservesExactThreshold() {
@@ -56,6 +58,7 @@ public void testBucketCollapsingPreservesExactThreshold() {
         assertThat(iterator.hasNext(), equalTo(false));
         assertThat(merged.zeroThreshold(), equalTo(3.0));
         assertThat(merged.count(), equalTo(52L));
+        assertThat(merged.isIndexBased(), equalTo(false));
     }
 
     public void testHashCodeEquality() {

server/build.gradle

@@ -38,6 +38,7 @@ dependencies {
   api project(":libs:plugin-analysis-api")
   api project(':libs:grok')
   api project(":libs:tdigest")
+  api project(":libs:exponential-histogram")
   implementation project(":libs:simdvec")
   implementation project(":libs:entitlement")
 

server/src/main/java/module-info.java

@@ -30,6 +30,7 @@
     requires org.elasticsearch.plugin.analysis;
     requires org.elasticsearch.grok;
     requires org.elasticsearch.tdigest;
+    requires org.elasticsearch.exponentialhistogram;
     requires org.elasticsearch.simdvec;
     requires org.elasticsearch.entitlement;