diff --git a/server/src/main/java/org/elasticsearch/search/query/TopDocsCollectorContext.java b/server/src/main/java/org/elasticsearch/search/query/TopDocsCollectorContext.java index 6312fd3856b..2314d11e7e3 100644 --- a/server/src/main/java/org/elasticsearch/search/query/TopDocsCollectorContext.java +++ b/server/src/main/java/org/elasticsearch/search/query/TopDocsCollectorContext.java @@ -35,6 +35,7 @@ import org.apache.lucene.search.TermQuery; import org.apache.lucene.search.TopDocs; import org.apache.lucene.search.TopDocsCollector; import org.apache.lucene.search.TopFieldCollector; +import org.apache.lucene.search.TopFieldDocs; import org.apache.lucene.search.TopScoreDocCollector; import org.apache.lucene.search.TotalHitCountCollector; import org.apache.lucene.search.TotalHits; @@ -189,11 +190,11 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext { } } - private final Collector collector; protected final @Nullable SortAndFormats sortAndFormats; - protected final Supplier totalHitsSupplier; - protected final Supplier topDocsSupplier; - protected final Supplier maxScoreSupplier; + private final Collector collector; + private final Supplier totalHitsSupplier; + private final Supplier topDocsSupplier; + private final Supplier maxScoreSupplier; /** * Ctr @@ -262,11 +263,23 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext { return collector; } + TopDocsAndMaxScore newTopDocs() { + TopDocs in = topDocsSupplier.get(); + float maxScore = maxScoreSupplier.get(); + final TopDocs newTopDocs; + if (in instanceof TopFieldDocs) { + TopFieldDocs fieldDocs = (TopFieldDocs) in; + newTopDocs = new TopFieldDocs(totalHitsSupplier.get(), fieldDocs.scoreDocs, fieldDocs.fields); + } else { + newTopDocs = new TopDocs(totalHitsSupplier.get(), in.scoreDocs); + } + return new TopDocsAndMaxScore(newTopDocs, maxScore); + } + @Override void postProcess(QuerySearchResult result) throws IOException { - final TopDocs topDocs = new TopDocs(totalHitsSupplier.get(), topDocsSupplier.get().scoreDocs); - result.topDocs(new TopDocsAndMaxScore(topDocs, maxScoreSupplier.get()), - sortAndFormats == null ? null : sortAndFormats.formats); + final TopDocsAndMaxScore topDocs = newTopDocs(); + result.topDocs(topDocs, sortAndFormats == null ? null : sortAndFormats.formats); } } @@ -291,27 +304,25 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext { @Override void postProcess(QuerySearchResult result) throws IOException { - final TopDocs topDocs = new TopDocs(totalHitsSupplier.get(), topDocsSupplier.get().scoreDocs); - final float maxScore; + final TopDocsAndMaxScore topDocs = newTopDocs(); if (scrollContext.totalHits == null) { // first round - scrollContext.totalHits = topDocs.totalHits; - maxScore = scrollContext.maxScore = maxScoreSupplier.get(); + scrollContext.totalHits = topDocs.topDocs.totalHits; + scrollContext.maxScore = topDocs.maxScore; } else { // subsequent round: the total number of hits and // the maximum score were computed on the first round - topDocs.totalHits = scrollContext.totalHits; - maxScore = scrollContext.maxScore; + topDocs.topDocs.totalHits = scrollContext.totalHits; + topDocs.maxScore = scrollContext.maxScore; } if (numberOfShards == 1) { // if we fetch the document in the same roundtrip, we already know the last emitted doc - if (topDocs.scoreDocs.length > 0) { + if (topDocs.topDocs.scoreDocs.length > 0) { // set the last emitted doc - scrollContext.lastEmittedDoc = topDocs.scoreDocs[topDocs.scoreDocs.length - 1]; + scrollContext.lastEmittedDoc = topDocs.topDocs.scoreDocs[topDocs.topDocs.scoreDocs.length - 1]; } } - result.topDocs(new TopDocsAndMaxScore(topDocs, maxScore), - sortAndFormats == null ? null : sortAndFormats.formats); + result.topDocs(topDocs, sortAndFormats == null ? null : sortAndFormats.formats); } }