Implement max for exponential histograms (#133913)

---------

Co-authored-by: elasticsearchmachine <[email protected]>

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, histogram.sum(), histogram.min(), data);
+            result.reset(histogram.zeroBucket().zeroThreshold(), totalCount, histogram.sum(), histogram.min(), histogram.max(), data);
             return result;
         } catch (IOException e) {
             throw new RuntimeException(e);

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

@@ -87,6 +87,11 @@ public double min() {
         return Double.NaN;
     }
 
+    @Override
+    public double max() {
+        return Double.NaN;
+    }
+
     @Override
     public long ramBytesUsed() {
         return 0;

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

@@ -47,7 +47,6 @@
  */
 public interface ExponentialHistogram extends Accountable {
 
-    // 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.
 
@@ -117,6 +116,13 @@ public interface ExponentialHistogram extends Accountable {
      */
     double min();
 
+    /**
+     * Returns maximum of all values represented by this histogram.
+     *
+     * @return the maximum, NaN for empty histograms
+     */
+    double max();
+
     /**
      * Represents a bucket range of an {@link ExponentialHistogram}, either the positive or the negative range.
      */
@@ -154,6 +160,7 @@ static boolean equals(ExponentialHistogram a, ExponentialHistogram b) {
         return a.scale() == b.scale()
             && a.sum() == b.sum()
             && equalsIncludingNaN(a.min(), b.min())
+            && equalsIncludingNaN(a.max(), b.max())
             && a.zeroBucket().equals(b.zeroBucket())
             && bucketIteratorsEqual(a.negativeBuckets().iterator(), b.negativeBuckets().iterator())
             && bucketIteratorsEqual(a.positiveBuckets().iterator(), b.positiveBuckets().iterator());
@@ -187,6 +194,7 @@ static int hashCode(ExponentialHistogram histogram) {
         hash = 31 * hash + Double.hashCode(histogram.sum());
         hash = 31 * hash + Long.hashCode(histogram.valueCount());
         hash = 31 * hash + Double.hashCode(histogram.min());
+        hash = 31 * hash + Double.hashCode(histogram.max());
         hash = 31 * hash + histogram.zeroBucket().hashCode();
         // we intentionally don't include the hash of the buckets here, because that is likely expensive to compute
         // instead, we assume that the value count and sum are a good enough approximation in most cases to minimize collisions

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

@@ -126,6 +126,7 @@ private void mergeValuesToHistogram() {
         Aggregates aggregates = rawValuesAggregates();
         valueBuffer.setSum(aggregates.sum());
         valueBuffer.setMin(aggregates.min());
+        valueBuffer.setMax(aggregates.max());
         int scale = valueBuffer.scale();
 
         // Buckets must be provided with their indices in ascending order.
@@ -166,15 +167,17 @@ private void mergeValuesToHistogram() {
 
     private Aggregates rawValuesAggregates() {
         if (valueCount == 0) {
-            return new Aggregates(0, Double.NaN);
+            return new Aggregates(0, Double.NaN, Double.NaN);
         }
         double sum = 0;
         double min = Double.MAX_VALUE;
+        double max = -Double.MAX_VALUE;
         for (int i = 0; i < valueCount; i++) {
             sum += rawValueBuffer[i];
             min = Math.min(min, rawValueBuffer[i]);
+            max = Math.max(max, rawValueBuffer[i]);
         }
-        return new Aggregates(sum, min);
+        return new Aggregates(sum, min, max);
     }
 
     private static long estimateBaseSize(int numBuckets) {
@@ -198,5 +201,5 @@ public void close() {
         }
     }
 
-    private record Aggregates(double sum, double min) {}
+    private record Aggregates(double sum, double min, double max) {}
 }

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

@@ -27,6 +27,7 @@
 import org.elasticsearch.core.Releasable;
 
 import java.util.OptionalLong;
+import java.util.function.DoubleBinaryOperator;
 
 import static org.elasticsearch.exponentialhistogram.ExponentialScaleUtils.getMaximumScaleIncrease;
 
@@ -151,7 +152,8 @@ public void add(ExponentialHistogram toAdd) {
         }
         buffer.setZeroBucket(zeroBucket);
         buffer.setSum(a.sum() + b.sum());
-        buffer.setMin(nanAwareMin(a.min(), b.min()));
+        buffer.setMin(nanAwareAggregate(a.min(), b.min(), Math::min));
+        buffer.setMax(nanAwareAggregate(a.max(), b.max(), Math::max));
         // We attempt to bring everything to the scale of A.
         // This might involve increasing the scale for B, which would increase its indices.
         // We need to ensure that we do not exceed MAX_INDEX / MIN_INDEX in this case.
@@ -231,14 +233,14 @@ private static int putBuckets(
         return overflowCount;
     }
 
-    private static double nanAwareMin(double a, double b) {
+    private static double nanAwareAggregate(double a, double b, DoubleBinaryOperator aggregator) {
         if (Double.isNaN(a)) {
             return b;
         }
         if (Double.isNaN(b)) {
             return a;
         }
-        return Math.min(a, b);
+        return aggregator.applyAsDouble(a, b);
     }
 
 }

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

@@ -67,7 +67,7 @@ public static double estimateSum(BucketIterator negativeBuckets, BucketIterator
      * Estimates the minimum value of the histogram based on the populated buckets.
      * The returned value is guaranteed to be less than or equal to the exact minimum value of the histogram values.
      * If the histogram is empty, an empty Optional is returned.
-     *
+     * <p>
      * Note that this method can return +-Infinity if the histogram bucket boundaries are not representable in a double.
      *
      * @param zeroBucket the zero bucket of the histogram
@@ -102,4 +102,40 @@ public static OptionalDouble estimateMin(
         }
         return OptionalDouble.empty();
     }
+
+    /**
+     * Estimates the maximum value of the histogram based on the populated buckets.
+     * The returned value is guaranteed to be greater than or equal to the exact maximum value of the histogram values.
+     * If the histogram is empty, an empty Optional is returned.
+     * <p>
+     * Note that this method can return +-Infinity if the histogram bucket boundaries are not representable in a double.
+     *
+     * @param zeroBucket the zero bucket of the histogram
+     * @param negativeBuckets the negative buckets of the histogram
+     * @param positiveBuckets the positive buckets of the histogram
+     * @return the estimated minimum
+     */
+    public static OptionalDouble estimateMax(
+        ZeroBucket zeroBucket,
+        ExponentialHistogram.Buckets negativeBuckets,
+        ExponentialHistogram.Buckets positiveBuckets
+    ) {
+        int scale = negativeBuckets.iterator().scale();
+        assert scale == positiveBuckets.iterator().scale();
+
+        OptionalLong positiveMaxIndex = positiveBuckets.maxBucketIndex();
+        if (positiveMaxIndex.isPresent()) {
+            return OptionalDouble.of(ExponentialScaleUtils.getUpperBucketBoundary(positiveMaxIndex.getAsLong(), scale));
+        }
+
+        if (zeroBucket.count() > 0) {
+            return OptionalDouble.of(zeroBucket.zeroThreshold());
+        }
+
+        BucketIterator negativeBucketsIt = negativeBuckets.iterator();
+        if (negativeBucketsIt.hasNext()) {
+            return OptionalDouble.of(-ExponentialScaleUtils.getLowerBucketBoundary(negativeBucketsIt.peekIndex(), scale));
+        }
+        return OptionalDouble.empty();
+    }
 }

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

@@ -33,6 +33,7 @@ public class ExponentialHistogramXContent {
     public static final String SCALE_FIELD = "scale";
     public static final String SUM_FIELD = "sum";
     public static final String MIN_FIELD = "min";
+    public static final String MAX_FIELD = "max";
     public static final String ZERO_FIELD = "zero";
     public static final String ZERO_COUNT_FIELD = "count";
     public static final String ZERO_THRESHOLD_FIELD = "threshold";
@@ -55,6 +56,9 @@ public static void serialize(XContentBuilder builder, ExponentialHistogram histo
         if (Double.isNaN(histogram.min()) == false) {
             builder.field(MIN_FIELD, histogram.min());
         }
+        if (Double.isNaN(histogram.max()) == false) {
+            builder.field(MAX_FIELD, histogram.max());
+        }
         double zeroThreshold = histogram.zeroBucket().zeroThreshold();
         long zeroCount = histogram.zeroBucket().count();
 

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

@@ -55,6 +55,7 @@ final class FixedCapacityExponentialHistogram extends AbstractExponentialHistogr
 
     private double sum;
     private double min;
+    private double max;
 
     private final ExponentialHistogramCircuitBreaker circuitBreaker;
     private boolean closed = false;
@@ -83,6 +84,7 @@ private FixedCapacityExponentialHistogram(int bucketCapacity, ExponentialHistogr
     void reset() {
         sum = 0;
         min = Double.NaN;
+        max = Double.NaN;
         setZeroBucket(ZeroBucket.minimalEmpty());
         resetBuckets(MAX_SCALE);
     }
@@ -133,6 +135,15 @@ void setMin(double min) {
         this.min = min;
     }
 
+    @Override
+    public double max() {
+        return max;
+    }
+
+    void setMax(double max) {
+        this.max = max;
+    }
+
     /**
      * Attempts to add a bucket to the positive or negative range of this histogram.
      * <br>

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

@@ -41,6 +41,7 @@ public enum Modification {
         SCALE,
         SUM,
         MIN,
+        MAX,
         ZERO_THRESHOLD,
         ZERO_COUNT,
         POSITIVE_BUCKETS,
@@ -102,6 +103,11 @@ private ExponentialHistogram copyWithModification(ExponentialHistogram toCopy, M
         } else {
             copy.setMin(toCopy.min());
         }
+        if (modification == Modification.MAX) {
+            copy.setMax(randomDouble());
+        } else {
+            copy.setMax(toCopy.max());
+        }
         long zeroCount = toCopy.zeroBucket().count();
         double zeroThreshold = toCopy.zeroBucket().zeroThreshold();
         if (modification == Modification.ZERO_COUNT) {

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

@@ -112,6 +112,7 @@ public void testAggregatesCorrectness() {
         double[] secondValues = randomDoubles(50).map(val -> val * 2 - 1).toArray();
         double correctSum = Arrays.stream(firstValues).sum() + Arrays.stream(secondValues).sum();
         double correctMin = DoubleStream.concat(Arrays.stream(firstValues), Arrays.stream(secondValues)).min().getAsDouble();
+        double correctMax = DoubleStream.concat(Arrays.stream(firstValues), Arrays.stream(secondValues)).max().getAsDouble();
         try (
             // Merge some empty histograms too to test that code path
             ReleasableExponentialHistogram merged = ExponentialHistogram.merge(
@@ -125,6 +126,7 @@ public void testAggregatesCorrectness() {
         ) {
             assertThat(merged.sum(), closeTo(correctSum, 0.000001));
             assertThat(merged.min(), equalTo(correctMin));
+            assertThat(merged.max(), equalTo(correctMax));
         }
     }