Add sum to exponential histograms (#133381)

Author: JonasKunz

benchmarks/src/main/java/org/elasticsearch/benchmark/exponentialhistogram/ExponentialHistogramMergeBench.java

@@ -130,7 +130,7 @@ private ExponentialHistogram asCompressedHistogram(ExponentialHistogram histogra
             CompressedExponentialHistogram.writeHistogramBytes(histoBytes, histogram.scale(), negativeBuckets, positiveBuckets);
             CompressedExponentialHistogram result = new CompressedExponentialHistogram();
             BytesRef data = histoBytes.bytes().toBytesRef();
-            result.reset(histogram.zeroBucket().zeroThreshold(), totalCount, data);
+            result.reset(histogram.zeroBucket().zeroThreshold(), totalCount, histogram.sum(), data);
             return result;
         } catch (IOException e) {
             throw new RuntimeException(e);

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

@@ -77,6 +77,11 @@ public Buckets negativeBuckets() {
         return EmptyBuckets.INSTANCE;
     }
 
+    @Override
+    public double sum() {
+        return 0;
+    }
+
     @Override
     public long ramBytesUsed() {
         return 0;

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

@@ -47,7 +47,7 @@
  */
 public interface ExponentialHistogram extends Accountable {
 
-    // TODO(b/128622): support min/max/sum/count storage and merging.
+    // TODO(b/128622): support min/max storage and merging.
     // TODO(b/128622): Add special positive and negative infinity buckets
     // to allow representation of explicit bucket histograms with open boundaries.
 
@@ -93,6 +93,15 @@ public interface ExponentialHistogram extends Accountable {
      */
     Buckets negativeBuckets();
 
+    /**
+     * Returns the sum of all values represented by this histogram.
+     * Note that even if histograms are cumulative, the sum is not guaranteed to be monotonically increasing,
+     * because histograms support negative values.
+     *
+     * @return the sum, guaranteed to be zero for empty histograms
+     */
+    double sum();
+
     /**
      * Represents a bucket range of an {@link ExponentialHistogram}, either the positive or the negative range.
      */

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

@@ -123,6 +123,7 @@ private void mergeValuesToHistogram() {
         }
 
         valueBuffer.reset();
+        valueBuffer.setSum(rawValuesSum());
         int scale = valueBuffer.scale();
 
         // Buckets must be provided with their indices in ascending order.
@@ -161,6 +162,14 @@ private void mergeValuesToHistogram() {
         valueCount = 0;
     }
 
+    private double rawValuesSum() {
+        double sum = 0;
+        for (int i = 0; i < valueCount; i++) {
+            sum += rawValueBuffer[i];
+        }
+        return sum;
+    }
+
     private static long estimateBaseSize(int numBuckets) {
         return SHALLOW_SIZE + RamEstimationUtil.estimateDoubleArray(numBuckets);
     };

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

@@ -150,6 +150,7 @@ public void add(ExponentialHistogram toAdd) {
             buffer = FixedCapacityExponentialHistogram.create(bucketLimit, circuitBreaker);
         }
         buffer.setZeroBucket(zeroBucket);
+        buffer.setSum(a.sum() + b.sum());
 
         // We attempt to bring everything to the scale of A.
         // This might involve increasing the scale for B, which would increase its indices.

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

@@ -0,0 +1,62 @@
+/*
+ * 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;
+
+public class ExponentialHistogramUtils {
+
+    /**
+     * Estimates the sum of all values of a histogram just based on the populated buckets.
+     * Will never return NaN, but might return +/-Infinity if the histogram is too big.
+     *
+     * @param negativeBuckets the negative buckets of the histogram
+     * @param positiveBuckets the positive buckets of the histogram
+     * @return the estimated sum of all values in the histogram, guaranteed to be zero if there are no buckets.
+     */
+    public static double estimateSum(BucketIterator negativeBuckets, BucketIterator positiveBuckets) {
+        assert negativeBuckets.scale() == positiveBuckets.scale();
+
+        // for each bucket index, sum up the counts, but account for the positive/negative sign
+        BucketIterator it = new MergingBucketIterator(
+            positiveBuckets,
+            negativeBuckets,
+            positiveBuckets.scale(),
+            (positiveCount, negativeCount) -> positiveCount - negativeCount
+        );
+        double sum = 0.0;
+        while (it.hasNext()) {
+            long countWithSign = it.peekCount();
+            double bucketMidPoint = ExponentialScaleUtils.getPointOfLeastRelativeError(it.peekIndex(), it.scale());
+            if (countWithSign != 0) { // avoid 0 * INFINITY = NaN
+                double toAdd = bucketMidPoint * countWithSign;
+                if (Double.isFinite(toAdd)) {
+                    sum += toAdd;
+                } else {
+                    // Avoid NaN in case we end up with e.g. -Infinity+Infinity
+                    // we consider the bucket with the bigger index the winner for the sign
+                    sum = toAdd;
+                }
+            }
+            it.advance();
+        }
+        return sum;
+    }
+}

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

@@ -31,6 +31,7 @@
 public class ExponentialHistogramXContent {
 
     public static final String SCALE_FIELD = "scale";
+    public static final String SUM_FIELD = "sum";
     public static final String ZERO_FIELD = "zero";
     public static final String ZERO_COUNT_FIELD = "count";
     public static final String ZERO_THRESHOLD_FIELD = "threshold";
@@ -49,6 +50,7 @@ public static void serialize(XContentBuilder builder, ExponentialHistogram histo
         builder.startObject();
 
         builder.field(SCALE_FIELD, histogram.scale());
+        builder.field(SUM_FIELD, histogram.sum());
         double zeroThreshold = histogram.zeroBucket().zeroThreshold();
         long zeroCount = histogram.zeroBucket().count();
 

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

@@ -53,6 +53,8 @@ final class FixedCapacityExponentialHistogram implements ReleasableExponentialHi
 
     private final Buckets positiveBuckets = new Buckets(true);
 
+    private double sum;
+
     private final ExponentialHistogramCircuitBreaker circuitBreaker;
     private boolean closed = false;
 
@@ -78,6 +80,7 @@ private FixedCapacityExponentialHistogram(int bucketCapacity, ExponentialHistogr
      * Resets this histogram to the same state as a newly constructed one with the same capacity.
      */
     void reset() {
+        sum = 0;
         setZeroBucket(ZeroBucket.minimalEmpty());
         resetBuckets(MAX_SCALE);
     }
@@ -110,6 +113,15 @@ void setZeroBucket(ZeroBucket zeroBucket) {
         this.zeroBucket = zeroBucket;
     }
 
+    @Override
+    public double sum() {
+        return sum;
+    }
+
+    void setSum(double sum) {
+        this.sum = sum;
+    }
+
     /**
      * Attempts to add a bucket to the positive or negative range of this histogram.
      * <br>

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

@@ -21,6 +21,8 @@
 
 package org.elasticsearch.exponentialhistogram;
 
+import java.util.function.LongBinaryOperator;
+
 /**
  * An iterator that merges two bucket iterators, aligning them to a common scale and combining buckets with the same index.
  */
@@ -33,6 +35,8 @@ final class MergingBucketIterator implements BucketIterator {
     private long currentIndex;
     private long currentCount;
 
+    private final LongBinaryOperator countMergeOperator;
+
     /**
      * Creates a new merging iterator.
      *
@@ -41,8 +45,22 @@ final class MergingBucketIterator implements BucketIterator {
      * @param targetScale the histogram scale to which both iterators should be aligned
      */
     MergingBucketIterator(BucketIterator itA, BucketIterator itB, int targetScale) {
+        this(itA, itB, targetScale, Long::sum);
+    }
+
+    /**
+     * Creates a new merging iterator, using the provided operator to merge the counts.
+     * Note that the resulting count can be negative if the operator produces negative results.
+     *
+     * @param itA         the first iterator to merge
+     * @param itB         the second iterator to merge
+     * @param countMergeOperator the operator to use to merge counts of buckets with the same index
+     * @param targetScale the histogram scale to which both iterators should be aligned
+     */
+    MergingBucketIterator(BucketIterator itA, BucketIterator itB, int targetScale, LongBinaryOperator countMergeOperator) {
         this.itA = new ScaleAdjustingBucketIterator(itA, targetScale);
         this.itB = new ScaleAdjustingBucketIterator(itB, targetScale);
+        this.countMergeOperator = countMergeOperator;
         endReached = false;
         advance();
     }
@@ -64,19 +82,21 @@ public void advance() {
             idxB = itB.peekIndex();
         }
 
-        currentCount = 0;
         boolean advanceA = hasNextA && (hasNextB == false || idxA <= idxB);
         boolean advanceB = hasNextB && (hasNextA == false || idxB <= idxA);
+        long countA = 0;
+        long countB = 0;
         if (advanceA) {
             currentIndex = idxA;
-            currentCount += itA.peekCount();
+            countA = itA.peekCount();
             itA.advance();
         }
         if (advanceB) {
             currentIndex = idxB;
-            currentCount += itB.peekCount();
+            countB = itB.peekCount();
             itB.advance();
         }
+        currentCount = countMergeOperator.applyAsLong(countA, countB);
     }
 
     @Override
@@ -106,4 +126,5 @@ private void assertEndNotReached() {
             throw new IllegalStateException("Iterator has no more buckets");
         }
     }
+
 }

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

@@ -106,6 +106,22 @@ public void testEmptyZeroBucketIgnored() {
         assertThat(posBuckets.hasNext(), equalTo(false));
     }
 
+    public void testSumCorrectness() {
+        double[] firstValues = randomDoubles(100).map(val -> val * 2 - 1).toArray();
+        double[] secondValues = randomDoubles(50).map(val -> val * 2 - 1).toArray();
+        double correctSum = Arrays.stream(firstValues).sum() + Arrays.stream(secondValues).sum();
+        try (
+            ReleasableExponentialHistogram merged = ExponentialHistogram.merge(
+                2,
+                breaker(),
+                createAutoReleasedHistogram(10, firstValues),
+                createAutoReleasedHistogram(20, secondValues)
+            )
+        ) {
+            assertThat(merged.sum(), closeTo(correctSum, 0.000001));
+        }
+    }
+
     public void testUpscalingDoesNotExceedIndexLimits() {
         for (int i = 0; i < 4; i++) {