1 files changed, 26 insertions, 12 deletions
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedBytesHashSamplerAggregator.java b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedBytesHashSamplerAggregator.java
index 33532a9f43..472881b008 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedBytesHashSamplerAggregator.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedBytesHashSamplerAggregator.java
@@ -26,6 +26,7 @@ import org.apache.lucene.search.DiversifiedTopDocsCollector.ScoreDocKey;
 import org.apache.lucene.search.TopDocsCollector;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.index.fielddata.AbstractNumericDocValues;
 import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
@@ -97,20 +98,33 @@ public class DiversifiedBytesHashSamplerAggregator extends SamplerAggregator {
                 } catch (IOException e) {
                     throw new ElasticsearchException("Error reading values", e);
                 }
-                return new NumericDocValues() {
-                    @Override
-                    public long get(int doc) {
+                return new AbstractNumericDocValues() {
 
-                        values.setDocument(doc);
-                        final int valuesCount = values.count();
-                        if (valuesCount > 1) {
-                            throw new IllegalArgumentException("Sample diversifying key must be a single valued-field");
-                        }
-                        if (valuesCount == 1) {
-                            final BytesRef bytes = values.valueAt(0);
-                            return bytes.hashCode();
+                    private int docID = -1;
+
+                    @Override
+                    public boolean advanceExact(int target) throws IOException {
+                        docID = target;
+                        if (values.advanceExact(target)) {
+                            if (values.docValueCount() > 1) {
+                                throw new IllegalArgumentException(
+                                        "Sample diversifying key must be a single valued-field");
+                            }
+                            return true;
+                        } else {
+                            return false;
                         }
-                        return 0;
+                    }
+
+                    @Override
+                    public int docID() {
+                        return docID;
+                    }
+
+                    @Override
+                    public long longValue() throws IOException {
+                        final BytesRef bytes = values.nextValue();
+                        return bytes.hashCode();
                     }
                 };
             }