1 files changed, 43 insertions, 16 deletions
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedOrdinalsSamplerAggregator.java b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedOrdinalsSamplerAggregator.java
index a5d027debe..74f7a16d62 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedOrdinalsSamplerAggregator.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedOrdinalsSamplerAggregator.java
@@ -22,11 +22,12 @@ package org.elasticsearch.search.aggregations.bucket.sampler;
 import org.apache.lucene.index.DocValues;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.NumericDocValues;
-import org.apache.lucene.index.RandomAccessOrds;
 import org.apache.lucene.index.SortedDocValues;
+import org.apache.lucene.index.SortedSetDocValues;
 import org.apache.lucene.search.DiversifiedTopDocsCollector;
 import org.apache.lucene.search.DiversifiedTopDocsCollector.ScoreDocKey;
 import org.apache.lucene.search.TopDocsCollector;
+import org.elasticsearch.index.fielddata.AbstractNumericDocValues;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.aggregations.bucket.DeferringBucketCollector;
@@ -84,29 +85,55 @@ public class DiversifiedOrdinalsSamplerAggregator extends SamplerAggregator {
 
             @Override
             protected NumericDocValues getKeys(LeafReaderContext context) {
-                final RandomAccessOrds globalOrds = valuesSource.globalOrdinalsValues(context);
+                final SortedSetDocValues globalOrds = valuesSource.globalOrdinalsValues(context);
                 final SortedDocValues singleValues = DocValues.unwrapSingleton(globalOrds);
                 if (singleValues != null) {
-                    return new NumericDocValues() {
+                    return new AbstractNumericDocValues() {
+
+                        @Override
+                        public boolean advanceExact(int target) throws IOException {
+                            return singleValues.advanceExact(target);
+                        }
+
                         @Override
-                        public long get(int doc) {
-                            return singleValues.getOrd(doc);
+                        public int docID() {
+                            return singleValues.docID();
+                        }
+
+                        @Override
+                        public long longValue() throws IOException {
+                            return singleValues.ordValue();
                         }
                     };
                 }
-                return new NumericDocValues() {
+                return new AbstractNumericDocValues() {
+
+                    long value;
+
                     @Override
-                    public long get(int doc) {
-                        globalOrds.setDocument(doc);
-                        final long valuesCount = globalOrds.cardinality();
-                        if (valuesCount > 1) {
-                            throw new IllegalArgumentException("Sample diversifying key must be a single valued-field");
-                        }
-                        if (valuesCount == 1) {
-                            long result = globalOrds.ordAt(0);
-                            return result;
+                    public boolean advanceExact(int target) throws IOException {
+                        if (globalOrds.advanceExact(target)) {
+                            value = globalOrds.nextOrd();
+                            // Check there isn't a second value for this
+                            // document
+                            if (globalOrds.nextOrd() != SortedSetDocValues.NO_MORE_ORDS) {
+                                throw new IllegalArgumentException(
+                                        "Sample diversifying key must be a single valued-field");
+                            }
+                            return true;
+                        } else {
+                            return false;
                         }
-                        return Long.MIN_VALUE;
+                    }
+
+                    @Override
+                    public int docID() {
+                        return globalOrds.docID();
+                    }
+
+                    @Override
+                    public long longValue() throws IOException {
+                        return value;
                     }
                 };