Properly calculate sum of NaN and infinities

Author: liketic

core/src/main/java/org/elasticsearch/search/aggregations/metrics/avg/AvgAggregator.java

@@ -89,10 +89,17 @@ public void collect(int doc, long bucket) throws IOException {
                     double compensation = compensations.get(bucket);
 
                     for (int i = 0; i < valueCount; i++) {
-                        double corrected = values.nextValue() - compensation;
-                        double newSum = sum + corrected;
-                        compensation = (newSum - sum) - corrected;
-                        sum = newSum;
+                        double value = values.nextValue();
+                        if (Double.isNaN(value) || Double.isInfinite(value)) {
+                            sum += value;
+                            if (Double.isNaN(sum))
+                                break;
+                        } else if (Double.isFinite(sum)) {
+                            double corrected = value - compensation;
+                            double newSum = sum + corrected;
+                            compensation = (newSum - sum) - corrected;
+                            sum = newSum;
+                        }
                     }
                     sums.set(bucket, sum);
                     compensations.set(bucket, compensation);

core/src/main/java/org/elasticsearch/search/aggregations/metrics/avg/InternalAvg.java

@@ -95,11 +95,18 @@ public InternalAvg doReduce(List<InternalAggregation> aggregations, ReduceContex
         // Compute the sum of double values with Kahan summation algorithm which is more
         // accurate than naive summation.
         for (InternalAggregation aggregation : aggregations) {
-            count += ((InternalAvg) aggregation).count;
-            double corrected = ((InternalAvg) aggregation).sum - compensation;
-            double newSum = sum + corrected;
-            compensation = (newSum - sum) - corrected;
-            sum = newSum;
+            InternalAvg avg = (InternalAvg) aggregation;
+            count += avg.count;
+            if (Double.isNaN(sum) == false) {
+                if (Double.isNaN(avg.sum) || Double.isInfinite(avg.sum)) {
+                    sum += avg.sum;
+                } else if (Double.isFinite(sum)) {
+                    double corrected = avg.sum - compensation;
+                    double newSum = sum + corrected;
+                    compensation = (newSum - sum) - corrected;
+                    sum = newSum;
+                }
+            }
         }
         return new InternalAvg(getName(), sum, count, format, pipelineAggregators(), getMetaData());
     }

core/src/main/java/org/elasticsearch/search/aggregations/metrics/stats/InternalStats.java

@@ -160,10 +160,17 @@ public InternalStats doReduce(List<InternalAggregation> aggregations, ReduceCont
             max = Math.max(max, stats.getMax());
             // Compute the sum of double values with Kahan summation algorithm which is more
             // accurate than naive summation.
-            double corrected = stats.getSum() - compensation;
-            double newSum = sum + corrected;
-            compensation = (newSum - sum) - corrected;
-            sum = newSum;
+            if (Double.isNaN(sum) == false) {
+                double value = stats.getSum();
+                if (Double.isNaN(value) || Double.isInfinite(value)) {
+                    sum += value;
+                } else if (Double.isFinite(sum)) {
+                    double corrected = value - compensation;
+                    double newSum = sum + corrected;
+                    compensation = (newSum - sum) - corrected;
+                    sum = newSum;
+                }
+            }
         }
         return new InternalStats(name, count, sum, min, max, format, pipelineAggregators(), getMetaData());
     }

core/src/main/java/org/elasticsearch/search/aggregations/metrics/stats/StatsAggregator.java

@@ -110,10 +110,16 @@ public void collect(int doc, long bucket) throws IOException {
 
                     for (int i = 0; i < valuesCount; i++) {
                         double value = values.nextValue();
-                        double corrected = value - compensation;
-                        double newSum = sum + corrected;
-                        compensation = (newSum - sum) - corrected;
-                        sum = newSum;
+                        if (Double.isNaN(sum) == false) {
+                            if (Double.isNaN(value) || Double.isInfinite(value)) {
+                                sum += value;
+                            } else if (Double.isFinite(sum)) {
+                                double corrected = value - compensation;
+                                double newSum = sum + corrected;
+                                compensation = (newSum - sum) - corrected;
+                                sum = newSum;
+                            }
+                        }
                         min = Math.min(min, value);
                         max = Math.max(max, value);
                     }

core/src/main/java/org/elasticsearch/search/aggregations/metrics/stats/extended/ExtendedStatsAggregator.java

@@ -121,15 +121,23 @@ public void collect(int doc, long bucket) throws IOException {
                     double compensationOfSqr = compensationOfSqrs.get(bucket);
                     for (int i = 0; i < valuesCount; i++) {
                         double value = values.nextValue();
-                        double corrected = value - compensation;
-                        double newSum = sum + corrected;
-                        compensation = (newSum - sum) - corrected;
-                        sum = newSum;
-
-                        double correctedOfSqr = value * value - compensationOfSqr;
-                        double newSumOfSqr = sumOfSqr + correctedOfSqr;
-                        compensationOfSqr = (newSumOfSqr - sumOfSqr) - correctedOfSqr;
-                        sumOfSqr = newSumOfSqr;
+                        if (Double.isNaN(value) || Double.isInfinite(value)) {
+                            sum += value;
+                            sumOfSqr += value * value;
+                        } else {
+                            if (Double.isFinite(sum)) {
+                                double corrected = value - compensation;
+                                double newSum = sum + corrected;
+                                compensation = (newSum - sum) - corrected;
+                                sum = newSum;
+                            }
+                            if (Double.isFinite(sumOfSqr)) {
+                                double correctedOfSqr = value * value - compensationOfSqr;
+                                double newSumOfSqr = sumOfSqr + correctedOfSqr;
+                                compensationOfSqr = (newSumOfSqr - sumOfSqr) - correctedOfSqr;
+                                sumOfSqr = newSumOfSqr;
+                            }
+                        }
                         min = Math.min(min, value);
                         max = Math.max(max, value);
                     }

core/src/main/java/org/elasticsearch/search/aggregations/metrics/stats/extended/InternalExtendedStats.java

@@ -142,12 +142,23 @@ public String getStdDeviationBoundAsString(Bounds bound) {
     @Override
     public InternalExtendedStats doReduce(List<InternalAggregation> aggregations, ReduceContext reduceContext) {
         double sumOfSqrs = 0;
+        double compensationOfSqrs = 0;
         for (InternalAggregation aggregation : aggregations) {
             InternalExtendedStats stats = (InternalExtendedStats) aggregation;
             if (stats.sigma != sigma) {
                 throw new IllegalStateException("Cannot reduce other stats aggregations that have a different sigma");
             }
-            sumOfSqrs += stats.getSumOfSquares();
+            if (Double.isNaN(sumOfSqrs) == false) {
+                double value = stats.getSumOfSquares();
+                if (Double.isNaN(value) || Double.isInfinite(value)) {
+                    sumOfSqrs += value;
+                } else if (Double.isFinite(sumOfSqrs)) {
+                    double correctedOfSqrs = value - compensationOfSqrs;
+                    double newSumOfSqrs = sumOfSqrs + correctedOfSqrs;
+                    compensationOfSqrs = (newSumOfSqrs - sumOfSqrs) - correctedOfSqrs;
+                    sumOfSqrs = newSumOfSqrs;
+                }
+            }
         }
         final InternalStats stats = super.doReduce(aggregations, reduceContext);
         return new InternalExtendedStats(name, stats.getCount(), stats.getSum(), stats.getMin(), stats.getMax(), sumOfSqrs, sigma,

core/src/main/java/org/elasticsearch/search/aggregations/metrics/sum/InternalSum.java

@@ -24,7 +24,6 @@
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.aggregations.InternalAggregation;
 import org.elasticsearch.search.aggregations.metrics.InternalNumericMetricsAggregation;
-import org.elasticsearch.search.aggregations.metrics.avg.InternalAvg;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
 
 import java.io.IOException;
@@ -36,7 +35,7 @@ public class InternalSum extends InternalNumericMetricsAggregation.SingleValue i
     private final double sum;
 
     public InternalSum(String name, double sum, DocValueFormat formatter, List<PipelineAggregator> pipelineAggregators,
-            Map<String, Object> metaData) {
+                       Map<String, Object> metaData) {
         super(name, pipelineAggregators, metaData);
         this.sum = sum;
         this.format = formatter;
@@ -79,10 +78,17 @@ public InternalSum doReduce(List<InternalAggregation> aggregations, ReduceContex
         double sum = 0;
         double compensation = 0;
         for (InternalAggregation aggregation : aggregations) {
-            double corrected = ((InternalSum) aggregation).sum - compensation;
-            double newSum = sum + corrected;
-            compensation = (newSum - sum) - corrected;
-            sum = newSum;
+            double value = ((InternalSum) aggregation).sum;
+            if (Double.isNaN(value) || Double.isInfinite(value)) {
+                sum += value;
+                if (Double.isNaN(sum))
+                    break;
+            } else if (Double.isFinite(sum)) {
+                double corrected = value - compensation;
+                double newSum = sum + corrected;
+                compensation = (newSum - sum) - corrected;
+                sum = newSum;
+            }
         }
         return new InternalSum(name, sum, format, pipelineAggregators(), getMetaData());
     }

core/src/main/java/org/elasticsearch/search/aggregations/metrics/sum/SumAggregator.java

@@ -82,10 +82,17 @@ public void collect(int doc, long bucket) throws IOException {
                     double sum = sums.get(bucket);
                     double compensation = compensations.get(bucket);
                     for (int i = 0; i < valuesCount; i++) {
-                        double corrected = values.nextValue() - compensation;
-                        double newSum = sum + corrected;
-                        compensation = (newSum - sum) - corrected;
-                        sum = newSum;
+                        double value = values.nextValue();
+                        if (Double.isNaN(value) || Double.isInfinite(value)) {
+                            sum += value;
+                            if (Double.isNaN(sum))
+                                break;
+                        } else if (Double.isFinite(sum)) {
+                            double corrected = value - compensation;
+                            double newSum = sum + corrected;
+                            compensation = (newSum - sum) - corrected;
+                            sum = newSum;
+                        }
                     }
                     compensations.set(bucket, compensation);
                     sums.set(bucket, sum);

core/src/test/java/org/elasticsearch/search/aggregations/metrics/InternalExtendedStatsTests.java

@@ -21,13 +21,15 @@
 
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.search.DocValueFormat;
+import org.elasticsearch.search.aggregations.InternalAggregation;
 import org.elasticsearch.search.aggregations.ParsedAggregation;
 import org.elasticsearch.search.aggregations.metrics.stats.extended.ExtendedStats.Bounds;
 import org.elasticsearch.search.aggregations.metrics.stats.extended.InternalExtendedStats;
 import org.elasticsearch.search.aggregations.metrics.stats.extended.ParsedExtendedStats;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
 import org.elasticsearch.test.InternalAggregationTestCase;
 
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -188,4 +190,17 @@ protected InternalExtendedStats mutateInstance(InternalExtendedStats instance) {
         }
         return new InternalExtendedStats(name, count, sum, min, max, sumOfSqrs, sigma, formatter, pipelineAggregators, metaData);
     }
+
+    public void testSummationAccuracy() {
+        double[] values = new double[]{0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.9, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7};
+        double sigma = randomDouble();
+        List<InternalAggregation> aggregations = new ArrayList<>(values.length);
+        for (double sumOfSqrs : values) {
+            aggregations.add(new InternalExtendedStats("dummy1", 1,  0.0, 0.0, 0.0, sumOfSqrs, sigma, null, null, null));
+        }
+        InternalExtendedStats stats = new InternalExtendedStats("dummy", 1,  0.0, 0.0, 0.0, 0.0, sigma, null, null, null);
+        InternalExtendedStats reduced = stats.doReduce(aggregations, null);
+        assertEquals(13.5, reduced.getSumOfSquares(), 0d);
+        assertEquals("dummy", reduced.getName());
+    }
 }

core/src/test/java/org/elasticsearch/search/aggregations/metrics/InternalSumTests.java

@@ -49,15 +49,8 @@ protected Writeable.Reader<InternalSum> instanceReader() {
 
     @Override
     protected void assertReduced(InternalSum reduced, List<InternalSum> inputs) {
-        double expectedSum = 0;
-        double compensation = 0;
-        for (InternalSum aggregation : inputs) {
-            double corrected = aggregation.value() - compensation;
-            double newSum = expectedSum + corrected;
-            compensation = (newSum - expectedSum) - corrected;
-            expectedSum = newSum;
-        }
-        assertEquals(expectedSum, reduced.getValue(), 0.000d);
+        double expectedSum = inputs.stream().mapToDouble(InternalSum::getValue).sum();
+        assertEquals(expectedSum, reduced.getValue(), 0.0001d);
     }
 
     public void testSummationAccuracy() throws IOException {