From 586093ec5e96bf657bc89040b0b6c822fd4f4678 Mon Sep 17 00:00:00 2001
From: Jim Ferenczi <jimczi@apache.org>
Date: Thu, 10 Jan 2019 00:26:35 +0100
Subject: [PATCH] Handle TopFieldDocs copy in TopDocsCollectorContext

This commit fixes the clone of TopFieldDocs.

Relates #37179
Relates #37266
---
 .../search/query/TopDocsCollectorContext.java | 45 ++++++++++++-------
 1 file changed, 28 insertions(+), 17 deletions(-)

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<TotalHits> totalHitsSupplier;
-        protected final Supplier<TopDocs> topDocsSupplier;
-        protected final Supplier<Float> maxScoreSupplier;
+        private final Collector collector;
+        private final Supplier<TotalHits> totalHitsSupplier;
+        private final Supplier<TopDocs> topDocsSupplier;
+        private final Supplier<Float> 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);
         }
     }