pass down docids in order to determine whether to load sequential

Author: martijnvg

server/src/main/java/org/elasticsearch/index/engine/LuceneSyntheticSourceChangesSnapshot.java

@@ -9,6 +9,8 @@
 
 package org.elasticsearch.index.engine;
 
+import com.carrotsearch.hppc.IntArrayList;
+
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.FieldDoc;
 import org.apache.lucene.search.ScoreDoc;
@@ -84,14 +86,7 @@ public LuceneSyntheticSourceChangesSnapshot(
         this.sourceLoader = mapperService.mappingLookup().newSourceLoader(null, SourceFieldMetrics.NOOP);
         Set<String> storedFields = sourceLoader.requiredStoredFields();
 
-        // If more than a few ops are requested let's enforce a sequential reader. Typically, thousands of ops may be requested.
-        // Given how LuceneSyntheticSourceChangesSnapshot accesses stored field, it should always benefit from using sequential reader.
-        //
-        // A sequential reader decompresses a block eagerly, so that increasing adjacent doc ids can access stored fields without
-        // compressing on each StoredFields#document(docId) invocation. The only downside is the last few operations in the request
-        // seq_no range are at the beginning of a block, which means stored fields for many docs are being decompressed that isn't used.
-        boolean shouldForceSequentialReader = toSeqNo - fromSeqNo > 10;
-        this.storedFieldLoader = StoredFieldLoader.create(false, storedFields, shouldForceSequentialReader);
+        this.storedFieldLoader = StoredFieldLoader.createLoaderWithMaybeSequentialReader(storedFields);
         this.lastSeenSeqNo = fromSeqNo - 1;
     }
 
@@ -199,8 +194,19 @@ private Translog.Operation[] loadDocuments(List<SearchRecord> documentRecords) t
                     maxDoc = leafReaderContext.reader().maxDoc();
                 } while (docRecord.docID() >= docBase + maxDoc);
 
-                leafFieldLoader = storedFieldLoader.getLoader(leafReaderContext, null);
-                leafSourceLoader = sourceLoader.leaf(leafReaderContext.reader(), null);
+                IntArrayList nextDocIds = new IntArrayList();
+                for (int j = i; j < documentRecords.size(); j++) {
+                    int docID = documentRecords.get(j).docID();
+                    if (docBase + maxDoc >= docID) {
+                        break;
+                    }
+                    int segmentDocID = docID - docBase;
+                    nextDocIds.add(segmentDocID);
+                }
+
+                int[] nextDocIdArray = nextDocIds.toArray();
+                leafFieldLoader = storedFieldLoader.getLoader(leafReaderContext, nextDocIdArray);
+                leafSourceLoader = sourceLoader.leaf(leafReaderContext.reader(), nextDocIdArray);
                 setNextSourceMetadataReader(leafReaderContext);
             }
             int segmentDocID = docRecord.docID() - docBase;

server/src/main/java/org/elasticsearch/index/fieldvisitor/StoredFieldLoader.java

@@ -235,4 +235,35 @@ public Map<String, List<Object>> storedFields() {
         }
     }
 
+    public static StoredFieldLoader createLoaderWithMaybeSequentialReader(Set<String> fields) {
+        List<String> fieldsToLoad = fieldsToLoad(false, fields);
+        return new StoredFieldLoader() {
+            @Override
+            public LeafStoredFieldLoader getLoader(LeafReaderContext ctx, int[] docs) throws IOException {
+                return new ReaderStoredFieldLoader(maybeSequentialReaderWithSmallDocIdGaps(ctx, docs), false, fields);
+            }
+
+            @Override
+            public List<String> fieldsToLoad() {
+                return fieldsToLoad;
+            }
+        };
+    }
+
+    private static CheckedBiConsumer<Integer, FieldsVisitor, IOException> maybeSequentialReaderWithSmallDocIdGaps(
+        LeafReaderContext ctx,
+        int[] docs
+    ) throws IOException {
+        LeafReader leafReader = ctx.reader();
+        if (docs != null && docs.length > 10 && hasNearSequentialDocs(docs)) {
+            return sequentialReader(ctx);
+        }
+        StoredFields storedFields = leafReader.storedFields();
+        return storedFields::document;
+    }
+
+    private static boolean hasNearSequentialDocs(int[] docs) {
+        return docs.length > 0 && docs[docs.length - 1] - docs[0] <= (1.2 * docs.length - 1);
+    }
+
 }