Speedup a few more aggregations using Kahan summation (#122920)

original-brownbear · web-flow · commit cc29af3b14c1 · 2025-02-19T11:50:09.000+01:00
Follow-up to #120241 and linked issues, using the logic that doesn't require any mutable object indirection led to significant speedups for sum+avg and should be of similar if not more help in these as well.
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/AvgAggregator.java
@@ -58,7 +58,7 @@ protected LeafBucketCollector getLeafCollector(NumericDoubleValues values, final
             public void collect(int doc, long bucket) throws IOException {
                 if (values.advanceExact(doc)) {
                     maybeGrow(bucket);
-                    computeSum(bucket, values, sums, compensations);
+                    computeSum(bucket, values.doubleValue(), sums, compensations);
                     counts.increment(bucket, 1L);
                 }
             }
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/ExtendedStatsAggregator.java b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/ExtendedStatsAggregator.java
@@ -105,7 +105,6 @@ public void collect(int doc, long bucket) throws IOException {
 
     @Override
     protected LeafBucketCollector getLeafCollector(NumericDoubleValues values, final LeafBucketCollector sub) {
-        final CompensatedSum compensatedSum = new CompensatedSum(0, 0);
         return new LeafBucketCollectorBase(sub, values) {
 
             @Override
@@ -114,20 +113,9 @@ public void collect(int doc, long bucket) throws IOException {
                     maybeGrow(bucket);
                     final double value = values.doubleValue();
                     counts.increment(bucket, 1L);
-                    // Compute the sum and sum of squires for double values with Kahan summation algorithm
-                    // which is more accurate than naive summation.
-                    compensatedSum.reset(sums.get(bucket), compensations.get(bucket));
-                    compensatedSum.add(value);
-                    sums.set(bucket, compensatedSum.value());
-                    compensations.set(bucket, compensatedSum.delta());
-
-                    compensatedSum.reset(sumOfSqrs.get(bucket), compensationOfSqrs.get(bucket));
-                    compensatedSum.add(value * value);
-                    sumOfSqrs.set(bucket, compensatedSum.value());
-                    compensationOfSqrs.set(bucket, compensatedSum.delta());
-
-                    mins.set(bucket, Math.min(mins.get(bucket), value));
-                    maxes.set(bucket, Math.max(maxes.get(bucket), value));
+                    SumAggregator.computeSum(bucket, value, sums, compensations);
+                    SumAggregator.computeSum(bucket, value * value, sumOfSqrs, compensationOfSqrs);
+                    StatsAggregator.updateMinsAndMaxes(bucket, value, mins, maxes);
                 }
             }
 
@@ -138,13 +126,14 @@ private void maybeGrow(long bucket) {
         if (bucket >= counts.size()) {
             final long from = counts.size();
             final long overSize = BigArrays.overSize(bucket + 1);
-            counts = bigArrays().resize(counts, overSize);
-            sums = bigArrays().resize(sums, overSize);
-            compensations = bigArrays().resize(compensations, overSize);
-            mins = bigArrays().resize(mins, overSize);
-            maxes = bigArrays().resize(maxes, overSize);
-            sumOfSqrs = bigArrays().resize(sumOfSqrs, overSize);
-            compensationOfSqrs = bigArrays().resize(compensationOfSqrs, overSize);
+            var bigArrays = bigArrays();
+            counts = bigArrays.resize(counts, overSize);
+            sums = bigArrays.resize(sums, overSize);
+            compensations = bigArrays.resize(compensations, overSize);
+            mins = bigArrays.resize(mins, overSize);
+            maxes = bigArrays.resize(maxes, overSize);
+            sumOfSqrs = bigArrays.resize(sumOfSqrs, overSize);
+            compensationOfSqrs = bigArrays.resize(compensationOfSqrs, overSize);
             mins.fill(from, overSize, Double.POSITIVE_INFINITY);
             maxes.fill(from, overSize, Double.NEGATIVE_INFINITY);
         }
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregator.java b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregator.java
@@ -93,7 +93,6 @@ public void collect(int doc, long bucket) throws IOException {
     }
 
     private LeafBucketCollector getLeafCollector(GeoPointValues values, LeafBucketCollector sub) {
-        final CompensatedSum compensatedSum = new CompensatedSum(0, 0);
         return new LeafBucketCollectorBase(sub, values) {
             @Override
             public void collect(int doc, long bucket) throws IOException {
@@ -104,16 +103,8 @@ public void collect(int doc, long bucket) throws IOException {
                     // Compute the sum of double values with Kahan summation algorithm which is more
                     // accurate than naive summation.
                     final GeoPoint value = values.pointValue();
-                    // latitude
-                    compensatedSum.reset(latSum.get(bucket), latCompensations.get(bucket));
-                    compensatedSum.add(value.getLat());
-                    latSum.set(bucket, compensatedSum.value());
-                    latCompensations.set(bucket, compensatedSum.delta());
-                    // longitude
-                    compensatedSum.reset(lonSum.get(bucket), lonCompensations.get(bucket));
-                    compensatedSum.add(value.getLon());
-                    lonSum.set(bucket, compensatedSum.value());
-                    lonCompensations.set(bucket, compensatedSum.delta());
+                    SumAggregator.computeSum(bucket, value.getLat(), latSum, latCompensations);
+                    SumAggregator.computeSum(bucket, value.getLon(), lonSum, lonCompensations);
                 }
             }
         };
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/StatsAggregator.java b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/StatsAggregator.java
@@ -81,36 +81,43 @@ public void collect(int doc, long bucket) throws IOException {
 
     @Override
     public LeafBucketCollector getLeafCollector(NumericDoubleValues values, LeafBucketCollector sub) {
-        final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
         return new LeafBucketCollectorBase(sub, values) {
             @Override
             public void collect(int doc, long bucket) throws IOException {
                 if (values.advanceExact(doc)) {
                     maybeGrow(bucket);
                     counts.increment(bucket, 1L);
-                    // Compute the sum of double values with Kahan summation algorithm which is more
-                    // accurate than naive summation.
-                    kahanSummation.reset(sums.get(bucket), compensations.get(bucket));
                     double value = values.doubleValue();
-                    kahanSummation.add(value);
-                    sums.set(bucket, kahanSummation.value());
-                    compensations.set(bucket, kahanSummation.delta());
-                    mins.set(bucket, Math.min(mins.get(bucket), value));
-                    maxes.set(bucket, Math.max(maxes.get(bucket), value));
+                    SumAggregator.computeSum(bucket, value, sums, compensations);
+                    updateMinsAndMaxes(bucket, value, mins, maxes);
                 }
             }
         };
     }
 
+    static void updateMinsAndMaxes(long bucket, double value, DoubleArray mins, DoubleArray maxes) {
+        double min = mins.get(bucket);
+        double updated = Math.min(value, min);
+        if (updated != min) {
+            mins.set(bucket, updated);
+        }
+        double max = maxes.get(bucket);
+        updated = Math.max(value, max);
+        if (updated != max) {
+            maxes.set(bucket, updated);
+        }
+    }
+
     private void maybeGrow(long bucket) {
         if (bucket >= counts.size()) {
             final long from = counts.size();
             final long overSize = BigArrays.overSize(bucket + 1);
-            counts = bigArrays().resize(counts, overSize);
-            sums = bigArrays().resize(sums, overSize);
-            compensations = bigArrays().resize(compensations, overSize);
-            mins = bigArrays().resize(mins, overSize);
-            maxes = bigArrays().resize(maxes, overSize);
+            var bigArrays = bigArrays();
+            counts = bigArrays.resize(counts, overSize);
+            sums = bigArrays.resize(sums, overSize);
+            compensations = bigArrays.resize(compensations, overSize);
+            mins = bigArrays.resize(mins, overSize);
+            maxes = bigArrays.resize(maxes, overSize);
             mins.fill(from, overSize, Double.POSITIVE_INFINITY);
             maxes.fill(from, overSize, Double.NEGATIVE_INFINITY);
         }
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/SumAggregator.java b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/SumAggregator.java
@@ -97,16 +97,15 @@ protected LeafBucketCollector getLeafCollector(NumericDoubleValues values, final
             public void collect(int doc, long bucket) throws IOException {
                 if (values.advanceExact(doc)) {
                     maybeGrow(bucket);
-                    computeSum(bucket, values, sums, compensations);
+                    computeSum(bucket, values.doubleValue(), sums, compensations);
                 }
             }
         };
     }
 
-    static void computeSum(long bucket, NumericDoubleValues values, DoubleArray sums, DoubleArray compensations) throws IOException {
+    static void computeSum(long bucket, double added, DoubleArray sums, DoubleArray compensations) {
         // Compute the sum of double values with Kahan summation algorithm which is more
         // accurate than naive summation.
-        double added = values.doubleValue();
         double value = addIfNonOrInf(added, sums.get(bucket));
         if (Double.isFinite(value)) {
             double delta = compensations.get(bucket);
@@ -122,8 +121,7 @@ static void computeSum(long bucket, NumericDoubleValues values, DoubleArray sums
 
     protected final void maybeGrow(long bucket) {
         if (bucket >= sums.size()) {
-            var bigArrays = bigArrays();
-            doGrow(bucket, bigArrays);
+            doGrow(bucket, bigArrays());
         }
     }
 

Original file line number	Diff line number	Diff line change
`@@ -58,7 +58,7 @@ protected LeafBucketCollector getLeafCollector(NumericDoubleValues values, final`
`58`	`58`	`public void collect(int doc, long bucket) throws IOException {`
`59`	`59`	`if (values.advanceExact(doc)) {`
`60`	`60`	`maybeGrow(bucket);`
`61`		`- computeSum(bucket, values, sums, compensations);`
	`61`	`+ computeSum(bucket, values.doubleValue(), sums, compensations);`
`62`	`62`	`counts.increment(bucket, 1L);`
`63`	`63`	`}`
`64`	`64`	`}`
Original file line number	Diff line number	Diff line change
`@@ -97,16 +97,15 @@ protected LeafBucketCollector getLeafCollector(NumericDoubleValues values, final`
`97`	`97`	`public void collect(int doc, long bucket) throws IOException {`
`98`	`98`	`if (values.advanceExact(doc)) {`
`99`	`99`	`maybeGrow(bucket);`
`100`		`- computeSum(bucket, values, sums, compensations);`
	`100`	`+ computeSum(bucket, values.doubleValue(), sums, compensations);`
`101`	`101`	`}`
`102`	`102`	`}`
`103`	`103`	`};`
`104`	`104`	`}`
`105`	`105`
`106`		`- static void computeSum(long bucket, NumericDoubleValues values, DoubleArray sums, DoubleArray compensations) throws IOException {`
	`106`	`+ static void computeSum(long bucket, double added, DoubleArray sums, DoubleArray compensations) {`
`107`	`107`	`// Compute the sum of double values with Kahan summation algorithm which is more`
`108`	`108`	`// accurate than naive summation.`
`109`		`- double added = values.doubleValue();`
`110`	`109`	`double value = addIfNonOrInf(added, sums.get(bucket));`
`111`	`110`	`if (Double.isFinite(value)) {`
`112`	`111`	`double delta = compensations.get(bucket);`
`@@ -122,8 +121,7 @@ static void computeSum(long bucket, NumericDoubleValues values, DoubleArray sums`
`122`	`121`
`123`	`122`	`protected final void maybeGrow(long bucket) {`
`124`	`123`	`if (bucket >= sums.size()) {`
`125`		`- var bigArrays = bigArrays();`
`126`		`- doGrow(bucket, bigArrays);`
	`124`	`+ doGrow(bucket, bigArrays());`
`127`	`125`	`}`
`128`	`126`	`}`
`129`	`127`