From a02a6a10af82557c25b3ab31c518d47c7305ed4c Mon Sep 17 00:00:00 2001
From: Thomas Kappler <tkappler@gmail.com>
Date: Fri, 6 Oct 2017 07:27:01 -0700
Subject: [PATCH] Fix IndexOutOfBoundsException in histograms for NaN doubles
 (#26787) (#26856)

---
 .../bucket/histogram/InternalHistogram.java   |  7 +++---
 .../histogram/InternalHistogramTests.java     | 24 +++++++++++++++++++
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/InternalHistogram.java b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/InternalHistogram.java
index 21cd2347cc61b..aa94bb762596a 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/InternalHistogram.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/InternalHistogram.java
@@ -321,8 +321,9 @@ protected boolean lessThan(IteratorAndCurrent a, IteratorAndCurrent b) {
             do {
                 final IteratorAndCurrent top = pq.top();
 
-                if (top.current.key != key) {
-                    // the key changes, reduce what we already buffered and reset the buffer for current buckets
+                if (Double.compare(top.current.key, key) != 0) {
+                    // The key changes, reduce what we already buffered and reset the buffer for current buckets.
+                    // Using Double.compare instead of != to handle NaN correctly.
                     final Bucket reduced = currentBuckets.get(0).reduce(currentBuckets, reduceContext);
                     if (reduced.getDocCount() >= minDocCount || reduceContext.isFinalReduce() == false) {
                         reducedBuckets.add(reduced);
@@ -335,7 +336,7 @@ protected boolean lessThan(IteratorAndCurrent a, IteratorAndCurrent b) {
 
                 if (top.iterator.hasNext()) {
                     final Bucket next = top.iterator.next();
-                    assert next.key > top.current.key : "shards must return data sorted by key";
+                    assert Double.compare(next.key, top.current.key) > 0 : "shards must return data sorted by key";
                     top.current = next;
                     pq.updateTop();
                 } else {
diff --git a/core/src/test/java/org/elasticsearch/search/aggregations/bucket/histogram/InternalHistogramTests.java b/core/src/test/java/org/elasticsearch/search/aggregations/bucket/histogram/InternalHistogramTests.java
index af826a7d7900e..8c383e799fee5 100644
--- a/core/src/test/java/org/elasticsearch/search/aggregations/bucket/histogram/InternalHistogramTests.java
+++ b/core/src/test/java/org/elasticsearch/search/aggregations/bucket/histogram/InternalHistogramTests.java
@@ -23,12 +23,15 @@
 import org.elasticsearch.common.io.stream.Writeable.Reader;
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.aggregations.BucketOrder;
+import org.elasticsearch.search.aggregations.InternalAggregation;
+import org.elasticsearch.search.aggregations.InternalAggregation.ReduceContext;
 import org.elasticsearch.search.aggregations.InternalAggregations;
 import org.elasticsearch.search.aggregations.InternalMultiBucketAggregationTestCase;
 import org.elasticsearch.search.aggregations.ParsedMultiBucketAggregation;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -63,6 +66,27 @@ protected InternalHistogram createTestInstance(String name,
         return new InternalHistogram(name, buckets, order, 1, null, format, keyed, pipelineAggregators, metaData);
     }
 
+    // issue 26787
+    public void testHandlesNaN() {
+        InternalHistogram histogram = createTestInstance();
+        InternalHistogram histogram2 = createTestInstance();
+        List<InternalHistogram.Bucket> buckets = histogram.getBuckets();
+        if (buckets == null || buckets.isEmpty()) {
+            return;
+        }
+
+        // Set the key of one bucket to NaN. Must be the last bucket because NaN is greater than everything else.
+        List<InternalHistogram.Bucket> newBuckets = new ArrayList<>(buckets.size());
+        if (buckets.size() > 1) {
+            newBuckets.addAll(buckets.subList(0, buckets.size() - 1));
+        }
+        InternalHistogram.Bucket b = buckets.get(buckets.size() - 1);
+        newBuckets.add(new InternalHistogram.Bucket(Double.NaN, b.docCount, keyed, b.format, b.aggregations));
+        
+        InternalHistogram newHistogram = histogram.create(newBuckets);
+        newHistogram.doReduce(Arrays.asList(newHistogram, histogram2), new InternalAggregation.ReduceContext(null, null, false));
+    }
+
     @Override
     protected void assertReduced(InternalHistogram reduced, List<InternalHistogram> inputs) {
         Map<Double, Long> expectedCounts = new TreeMap<>();