Sort document by internal doc id in FetchPhase to better use LRU cache (#57273)

This change sorts the docIdsToLoad once instead of in each sub-phase.

server/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java

@@ -61,6 +61,7 @@ import org.elasticsearch.search.lookup.SourceLookup;
 import org.elasticsearch.tasks.TaskCancelledException;
 
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -144,13 +145,20 @@ public class FetchPhase implements SearchPhase {
         }
 
         try {
+            DocIdToIndex[] docs = new DocIdToIndex[context.docIdsToLoadSize()];
+            for (int index = 0; index < context.docIdsToLoadSize(); index++) {
+                docs[index] = new DocIdToIndex(context.docIdsToLoad()[context.docIdsToLoadFrom() + index], index);
+            }
+            Arrays.sort(docs);
+
             SearchHit[] hits = new SearchHit[context.docIdsToLoadSize()];
+            SearchHit[] sortedHits = new SearchHit[context.docIdsToLoadSize()];
             FetchSubPhase.HitContext hitContext = new FetchSubPhase.HitContext();
             for (int index = 0; index < context.docIdsToLoadSize(); index++) {
                 if (context.isCancelled()) {
                     throw new TaskCancelledException("cancelled");
                 }
-                int docId = context.docIdsToLoad()[context.docIdsToLoadFrom() + index];
+                int docId = docs[index].docId;
                 int readerIndex = ReaderUtil.subIndex(docId, context.searcher().getIndexReader().leaves());
                 LeafReaderContext subReaderContext = context.searcher().getIndexReader().leaves().get(readerIndex);
                 int subDocId = docId - subReaderContext.docBase;
@@ -165,7 +173,8 @@ public class FetchPhase implements SearchPhase {
                         storedToRequestedFields, subReaderContext);
                 }
 
-                hits[index] = searchHit;
+                sortedHits[index] = searchHit;
+                hits[docs[index].index] = searchHit;
                 hitContext.reset(searchHit, subReaderContext, subDocId, context.searcher());
                 for (FetchSubPhase fetchSubPhase : fetchSubPhases) {
                     fetchSubPhase.hitExecute(context, hitContext);
@@ -176,7 +185,7 @@ public class FetchPhase implements SearchPhase {
             }
 
             for (FetchSubPhase fetchSubPhase : fetchSubPhases) {
-                fetchSubPhase.hitsExecute(context, hits);
+                fetchSubPhase.hitsExecute(context, sortedHits);
                 if (context.isCancelled()) {
                     throw new TaskCancelledException("cancelled");
                 }
@@ -189,6 +198,21 @@ public class FetchPhase implements SearchPhase {
         }
     }
 
+    static class DocIdToIndex implements Comparable<DocIdToIndex> {
+        final int docId;
+        final int index;
+
+        DocIdToIndex(int docId, int index) {
+            this.docId = docId;
+            this.index = index;
+        }
+
+        @Override
+        public int compareTo(DocIdToIndex o) {
+            return Integer.compare(docId, o.docId);
+        }
+    }
+
     private int findRootDocumentIfNested(SearchContext context, LeafReaderContext subReaderContext, int subDocId) throws IOException {
         if (context.mapperService().hasNested()) {
             BitSet bits = context.bitsetFilterCache()

server/src/main/java/org/elasticsearch/search/fetch/FetchSubPhase.java

@@ -81,6 +81,8 @@ public interface FetchSubPhase {
      */
     default void hitExecute(SearchContext context, HitContext hitContext) throws IOException {}
 
-
+    /**
+     * Executes the hits level phase (note, hits are sorted by doc ids).
+     */
     default void hitsExecute(SearchContext context, SearchHit[] hits) throws IOException {}
 }

server/src/main/java/org/elasticsearch/search/fetch/subphase/FetchDocValuesPhase.java

@@ -40,9 +40,7 @@ import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collections;
-import java.util.Comparator;
 import java.util.List;
 import java.util.Objects;
 
@@ -78,9 +76,6 @@ public final class FetchDocValuesPhase implements FetchSubPhase {
             return;
         }
 
-        hits = hits.clone(); // don't modify the incoming hits
-        Arrays.sort(hits, Comparator.comparingInt(SearchHit::docId));
-
         if (context.docValuesContext().fields().stream()
                 .map(f -> f.format)
                 .filter(USE_DEFAULT_FORMAT::equals)

server/src/main/java/org/elasticsearch/search/fetch/subphase/FetchScorePhase.java

@@ -30,8 +30,6 @@ import org.elasticsearch.search.fetch.FetchSubPhase;
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
-import java.util.Arrays;
-import java.util.Comparator;
 import java.util.Iterator;
 
 public class FetchScorePhase implements FetchSubPhase {
@@ -44,9 +42,6 @@ public class FetchScorePhase implements FetchSubPhase {
             return;
         }
 
-        hits = hits.clone(); // don't modify the incoming hits
-        Arrays.sort(hits, Comparator.comparingInt(SearchHit::docId));
-
         final IndexSearcher searcher = context.searcher();
         final Weight weight = searcher.createWeight(searcher.rewrite(context.query()), ScoreMode.COMPLETE, 1);
         Iterator<LeafReaderContext> leafContextIterator = searcher.getIndexReader().leaves().iterator();

server/src/main/java/org/elasticsearch/search/fetch/subphase/FetchVersionPhase.java

@@ -28,8 +28,6 @@ import org.elasticsearch.search.fetch.FetchSubPhase;
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
-import java.util.Arrays;
-import java.util.Comparator;
 
 public final class FetchVersionPhase implements FetchSubPhase {
     @Override
@@ -39,9 +37,6 @@ public final class FetchVersionPhase implements FetchSubPhase {
             return;
         }
 
-        hits = hits.clone(); // don't modify the incoming hits
-        Arrays.sort(hits, Comparator.comparingInt(SearchHit::docId));
-
         int lastReaderId = -1;
         NumericDocValues versions = null;
         for (SearchHit hit : hits) {

server/src/main/java/org/elasticsearch/search/fetch/subphase/MatchedQueriesPhase.java

@@ -35,7 +35,6 @@ import org.elasticsearch.search.internal.SearchContext.Lifetime;
 
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -49,8 +48,6 @@ public final class MatchedQueriesPhase implements FetchSubPhase {
             context.parsedQuery() == null) {
             return;
         }
-        hits = hits.clone(); // don't modify the incoming hits
-        Arrays.sort(hits, (a, b) -> Integer.compare(a.docId(), b.docId()));
         @SuppressWarnings("unchecked")
         List<String>[] matchedQueries = new List[hits.length];
         for (int i = 0; i < matchedQueries.length; ++i) {

server/src/main/java/org/elasticsearch/search/fetch/subphase/ScriptFieldsPhase.java

@@ -30,10 +30,8 @@ import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.Comparator;
 import java.util.List;
 
 public final class ScriptFieldsPhase implements FetchSubPhase {
@@ -44,9 +42,6 @@ public final class ScriptFieldsPhase implements FetchSubPhase {
             return;
         }
 
-        hits = hits.clone(); // don't modify the incoming hits
-        Arrays.sort(hits, Comparator.comparingInt(SearchHit::docId));
-
         int lastReaderId = -1;
         FieldScript[] leafScripts = null;
         List<ScriptFieldsContext.ScriptField> scriptFields = context.scriptFields().fields();

server/src/main/java/org/elasticsearch/search/fetch/subphase/SeqNoPrimaryTermPhase.java

@@ -28,8 +28,6 @@ import org.elasticsearch.search.fetch.FetchSubPhase;
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
-import java.util.Arrays;
-import java.util.Comparator;
 
 public final class SeqNoPrimaryTermPhase implements FetchSubPhase {
     @Override
@@ -38,9 +36,6 @@ public final class SeqNoPrimaryTermPhase implements FetchSubPhase {
             return;
         }
 
-        hits = hits.clone(); // don't modify the incoming hits
-        Arrays.sort(hits, Comparator.comparingInt(SearchHit::docId));
-
         int lastReaderId = -1;
         NumericDocValues seqNoField = null;
         NumericDocValues primaryTermField = null;