Add builder for exponential histograms (#133967)

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

@@ -206,6 +206,16 @@ public interface ExponentialHistogram extends Accountable {
         return EmptyExponentialHistogram.INSTANCE;
     }
 
+    /**
+     * Create a builder for an exponential histogram with the given scale.
+     * @param scale the scale of the histogram to build
+     * @param breaker the circuit breaker to use
+     * @return a new builder
+     */
+    static ExponentialHistogramBuilder builder(int scale, ExponentialHistogramCircuitBreaker breaker) {
+        return new ExponentialHistogramBuilder(scale, breaker);
+    }
+
     /**
      * Creates a histogram representing the distribution of the given values with at most the given number of buckets.
      * If the given {@code maxBucketCount} is greater than or equal to the number of values, the resulting histogram will have a

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

@@ -0,0 +1,159 @@
+/*
+ * Copyright Elasticsearch B.V., and/or licensed to Elasticsearch B.V.
+ * under one or more 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.
+ *
+ * This file is based on a modification of https://github.com/open-telemetry/opentelemetry-java which is licensed under the Apache 2.0 License.
+ */
+
+package org.elasticsearch.exponentialhistogram;
+
+import org.elasticsearch.core.Releasables;
+
+import java.util.TreeMap;
+
+/**
+ * A builder for building a {@link ReleasableExponentialHistogram} directly from buckets.
+ * Note that this class is not optimized regarding memory allocations, so it is not intended for high-throughput usage.
+ */
+public class ExponentialHistogramBuilder {
+
+    private final ExponentialHistogramCircuitBreaker breaker;
+
+    private final int scale;
+    private ZeroBucket zeroBucket = ZeroBucket.minimalEmpty();
+    private Double sum;
+    private Double min;
+    private Double max;
+
+    private final TreeMap<Long, Long> negativeBuckets = new TreeMap<>();
+    private final TreeMap<Long, Long> positiveBuckets = new TreeMap<>();
+
+    ExponentialHistogramBuilder(int scale, ExponentialHistogramCircuitBreaker breaker) {
+        this.breaker = breaker;
+        this.scale = scale;
+    }
+
+    public ExponentialHistogramBuilder zeroBucket(ZeroBucket zeroBucket) {
+        this.zeroBucket = zeroBucket;
+        return this;
+    }
+
+    /**
+     * Sets the sum of the histogram values. If not set, the sum will be estimated from the buckets.
+     * @param sum the sum value
+     * @return the builder
+     */
+    public ExponentialHistogramBuilder sum(double sum) {
+        this.sum = sum;
+        return this;
+    }
+
+    /**
+     * Sets the min value of the histogram values. If not set, the min will be estimated from the buckets.
+     * @param min the min value
+     * @return the builder
+     */
+    public ExponentialHistogramBuilder min(double min) {
+        this.min = min;
+        return this;
+    }
+
+    /**
+     * Sets the max value of the histogram values. If not set, the max will be estimated from the buckets.
+     * @param max the max value
+     * @return the builder
+     */
+    public ExponentialHistogramBuilder max(double max) {
+        this.max = max;
+        return this;
+    }
+
+    /**
+     * Adds the given bucket to the positive buckets.
+     * Buckets may be added in arbitrary order, but each bucket can only be added once.
+     *
+     * @param index the index of the bucket
+     * @param count the count of the bucket, must be at least 1
+     * @return the builder
+     */
+    public ExponentialHistogramBuilder addPositiveBucket(long index, long count) {
+        if (count < 1) {
+            throw new IllegalArgumentException("Bucket count must be at least 1");
+        }
+        if (positiveBuckets.containsKey(index)) {
+            throw new IllegalArgumentException("Positive bucket already exists: " + index);
+        }
+        positiveBuckets.put(index, count);
+        return this;
+    }
+
+    /**
+     * Adds the given bucket to the negative buckets.
+     * Buckets may be added in arbitrary order, but each bucket can only be added once.
+     *
+     * @param index the index of the bucket
+     * @param count the count of the bucket, must be at least 1
+     * @return the builder
+     */
+    public ExponentialHistogramBuilder addNegativeBucket(long index, long count) {
+        if (count < 1) {
+            throw new IllegalArgumentException("Bucket count must be at least 1");
+        }
+        if (negativeBuckets.containsKey(index)) {
+            throw new IllegalArgumentException("Negative bucket already exists: " + index);
+        }
+        negativeBuckets.put(index, count);
+        return this;
+    }
+
+    public ReleasableExponentialHistogram build() {
+        FixedCapacityExponentialHistogram result = FixedCapacityExponentialHistogram.create(
+            negativeBuckets.size() + positiveBuckets.size(),
+            breaker
+        );
+        boolean success = false;
+        try {
+            result.resetBuckets(scale);
+            result.setZeroBucket(zeroBucket);
+            negativeBuckets.forEach((index, count) -> result.tryAddBucket(index, count, false));
+            positiveBuckets.forEach((index, count) -> result.tryAddBucket(index, count, true));
+
+            double sumVal = (sum != null)
+                ? sum
+                : ExponentialHistogramUtils.estimateSum(result.negativeBuckets().iterator(), result.positiveBuckets().iterator());
+            double minVal = (min != null)
+                ? min
+                : ExponentialHistogramUtils.estimateMin(zeroBucket, result.negativeBuckets(), result.positiveBuckets()).orElse(Double.NaN);
+            double maxVal = (max != null)
+                ? max
+                : ExponentialHistogramUtils.estimateMax(zeroBucket, result.negativeBuckets(), result.positiveBuckets()).orElse(Double.NaN);
+
+            result.setMin(minVal);
+            result.setMax(maxVal);
+            result.setSum(sumVal);
+
+            success = true;
+        } finally {
+            if (success == false) {
+                Releasables.close(result);
+            }
+        }
+
+        // Create histogram
+        return result;
+    }
+}

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

@@ -0,0 +1,95 @@
+/*
+ * Copyright Elasticsearch B.V., and/or licensed to Elasticsearch B.V.
+ * under one or more 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.
+ *
+ * This file is based on a modification of https://github.com/open-telemetry/opentelemetry-java which is licensed under the Apache 2.0 License.
+ */
+
+package org.elasticsearch.exponentialhistogram;
+
+import java.util.List;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class ExponentialHistogramBuilderTests extends ExponentialHistogramTestCase {
+
+    public void testBuildWithAllFieldsSet() {
+        ZeroBucket zeroBucket = ZeroBucket.create(1, 2);
+        ExponentialHistogramBuilder builder = ExponentialHistogram.builder(3, breaker())
+            .zeroBucket(zeroBucket)
+            .sum(100.0)
+            .min(1.0)
+            .max(50.0)
+            .addPositiveBucket(2, 10)
+            .addPositiveBucket(0, 1)
+            .addPositiveBucket(5, 2)
+            .addNegativeBucket(-2, 5)
+            .addNegativeBucket(1, 2);
+
+        try (ReleasableExponentialHistogram histogram = builder.build()) {
+            assertThat(histogram.scale(), equalTo(3));
+            assertThat(histogram.zeroBucket(), equalTo(zeroBucket));
+            assertThat(histogram.sum(), equalTo(100.0));
+            assertThat(histogram.min(), equalTo(1.0));
+            assertThat(histogram.max(), equalTo(50.0));
+            assertBuckets(histogram.positiveBuckets(), List.of(0L, 2L, 5L), List.of(1L, 10L, 2L));
+            assertBuckets(histogram.negativeBuckets(), List.of(-2L, 1L), List.of(5L, 2L));
+        }
+    }
+
+    public void testBuildWithEstimation() {
+        ExponentialHistogramBuilder builder = ExponentialHistogram.builder(0, breaker())
+            .addPositiveBucket(0, 1)
+            .addPositiveBucket(1, 1)
+            .addNegativeBucket(0, 4);
+
+        try (ReleasableExponentialHistogram histogram = builder.build()) {
+            assertThat(histogram.scale(), equalTo(0));
+            assertThat(histogram.zeroBucket(), equalTo(ZeroBucket.minimalEmpty()));
+            assertThat(histogram.sum(), equalTo(-1.3333333333333335));
+            assertThat(histogram.min(), equalTo(-2.0));
+            assertThat(histogram.max(), equalTo(4.0));
+            assertBuckets(histogram.positiveBuckets(), List.of(0L, 1L), List.of(1L, 1L));
+            assertBuckets(histogram.negativeBuckets(), List.of(0L), List.of(4L));
+        }
+    }
+
+    public void testAddDuplicatePositiveBucketThrows() {
+        ExponentialHistogramBuilder builder = ExponentialHistogram.builder(0, breaker());
+        builder.addPositiveBucket(1, 10);
+        expectThrows(IllegalArgumentException.class, () -> builder.addPositiveBucket(1, 5));
+    }
+
+    public void testAddDuplicateNegativeBucketThrows() {
+        ExponentialHistogramBuilder builder = ExponentialHistogram.builder(0, breaker());
+        builder.addNegativeBucket(-1, 10);
+        expectThrows(IllegalArgumentException.class, () -> builder.addNegativeBucket(-1, 5));
+    }
+
+    private static void assertBuckets(ExponentialHistogram.Buckets buckets, List<Long> indices, List<Long> counts) {
+        List<Long> actualIndices = new java.util.ArrayList<>();
+        List<Long> actualCounts = new java.util.ArrayList<>();
+        BucketIterator it = buckets.iterator();
+        while (it.hasNext()) {
+            actualIndices.add(it.peekIndex());
+            actualCounts.add(it.peekCount());
+            it.advance();
+        }
+        assertThat("Expected bucket indices to match", actualIndices, equalTo(indices));
+        assertThat("Expected bucket counts to match", actualCounts, equalTo(counts));
+    }
+}