For the fields fetch phase, avoid reloading stored fields. (#58196)

This PR updates FetchFieldsPhase to override hitExecute instead of hitsExecute
(plural). This way, we can make sure that the stored fields (including _source)
are only loaded once per hit as part of FetchPhase.

Author: jtibshirani

rest-api-spec/src/main/resources/rest-api-spec/test/search/330_fetch_fields.yml

@@ -169,3 +169,48 @@ setup:
 
   - match: { hits.hits.0.fields.integer.0: 42 }
   - is_false: hits.hits.1.fields.integer
+
+---
+"Test disable _source loading":
+  - do:
+      indices.create:
+        index:  test
+        body:
+          settings:
+            number_of_shards: 1
+          mappings:
+            properties:
+              keyword:
+                type: keyword
+              integer:
+                type: integer
+                store: true
+
+  - do:
+      index:
+        index:  test
+        id:     1
+        refresh: true
+        body:
+          keyword: "x"
+          integer: 42
+
+  - do:
+      search:
+        index: test
+        body:
+          fields: [ keyword ]
+          _source: false
+
+  - match: { hits.hits.0.fields.keyword.0: "x" }
+
+  - do:
+      search:
+        index: test
+        body:
+          fields: [ keyword ]
+          stored_fields: [ integer ]
+          _source: false
+
+  - match: { hits.hits.0.fields.keyword.0: "x" }
+  - match: { hits.hits.0.fields.integer.0: 42 }

server/src/main/java/org/elasticsearch/search/SearchService.java

@@ -937,7 +937,10 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
             context.docValuesContext(new FetchDocValuesContext(docValueFields));
         }
         if (source.fetchFields() != null) {
-            context.fetchFieldsContext(new FetchFieldsContext(source.fetchFields()));
+            String indexName = context.indexShard().shardId().getIndexName();
+            FetchFieldsContext fetchFieldsContext = FetchFieldsContext.create(
+                indexName, context.mapperService(), source.fetchFields());
+            context.fetchFieldsContext(fetchFieldsContext);
         }
         if (source.highlighter() != null) {
             HighlightBuilder highlightBuilder = source.highlighter();

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

@@ -102,7 +102,8 @@ public void execute(SearchContext context) {
             if (!context.hasScriptFields() && !context.hasFetchSourceContext()) {
                 context.fetchSourceContext(new FetchSourceContext(true));
             }
-            fieldsVisitor = new FieldsVisitor(context.sourceRequested());
+            boolean loadSource = context.sourceRequested() || context.fetchFieldsContext() != null;
+            fieldsVisitor = new FieldsVisitor(loadSource);
         } else if (storedFieldsContext.fetchFields() == false) {
             // disable stored fields entirely
             fieldsVisitor = null;
@@ -131,7 +132,7 @@ public void execute(SearchContext context) {
                     }
                 }
             }
-            boolean loadSource = context.sourceRequested();
+            boolean loadSource = context.sourceRequested() || context.fetchFieldsContext() != null;
             if (storedToRequestedFields.isEmpty()) {
                 // empty list specified, default to disable _source if no explicit indication
                 fieldsVisitor = new FieldsVisitor(loadSource);

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

@@ -18,20 +18,36 @@
  */
 package org.elasticsearch.search.fetch.subphase;
 
+import org.elasticsearch.index.mapper.DocumentMapper;
+import org.elasticsearch.index.mapper.MapperService;
+
 import java.util.List;
 
 /**
  * The context needed to retrieve fields.
  */
 public class FetchFieldsContext {
 
-    private final List<FieldAndFormat> fields;
+    private FieldValueRetriever fieldValueRetriever;
+
+    public static FetchFieldsContext create(String indexName,
+                                            MapperService mapperService,
+                                            List<FieldAndFormat> fields) {
+        DocumentMapper documentMapper = mapperService.documentMapper();
+        if (documentMapper.sourceMapper().enabled() == false) {
+            throw new IllegalArgumentException("Unable to retrieve the requested [fields] since _source is " +
+                "disabled in the mappings for index [" + indexName + "]");
+        }
+
+        FieldValueRetriever fieldValueRetriever = FieldValueRetriever.create(mapperService, fields);
+        return new FetchFieldsContext(fieldValueRetriever);
+    }
 
-    public FetchFieldsContext(List<FieldAndFormat> fields) {
-        this.fields = fields;
+    private FetchFieldsContext(FieldValueRetriever fieldValueRetriever) {
+        this.fieldValueRetriever = fieldValueRetriever;
     }
 
-    public List<FieldAndFormat> fields() {
-        return this.fields;
+    public FieldValueRetriever fieldValueRetriever() {
+        return fieldValueRetriever;
     }
 }

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

@@ -19,10 +19,7 @@
 
 package org.elasticsearch.search.fetch.subphase;
 
-import org.apache.lucene.index.LeafReaderContext;
-import org.apache.lucene.index.ReaderUtil;
 import org.elasticsearch.common.document.DocumentField;
-import org.elasticsearch.index.mapper.DocumentMapper;
 import org.elasticsearch.index.mapper.IgnoredFieldMapper;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.fetch.FetchSubPhase;
@@ -40,33 +37,20 @@
 public final class FetchFieldsPhase implements FetchSubPhase {
 
     @Override
-    public void hitsExecute(SearchContext context, SearchHit[] hits) {
+    public void hitExecute(SearchContext context, HitContext hitContext) {
         FetchFieldsContext fetchFieldsContext = context.fetchFieldsContext();
-        if (fetchFieldsContext == null || fetchFieldsContext.fields().isEmpty()) {
+        if (fetchFieldsContext == null) {
             return;
         }
 
-        DocumentMapper documentMapper = context.mapperService().documentMapper();
-        if (documentMapper.sourceMapper().enabled() == false) {
-            throw new IllegalArgumentException("Unable to retrieve the requested [fields] since _source is " +
-                "disabled in the mappings for index [" + context.indexShard().shardId().getIndexName() + "]");
-        }
-
+        SearchHit hit = hitContext.hit();
         SourceLookup sourceLookup = context.lookup().source();
-        FieldValueRetriever fieldValueRetriever = FieldValueRetriever.create(
-            context.mapperService(),
-            fetchFieldsContext.fields());
-
-        for (SearchHit hit : hits) {
-            int readerIndex = ReaderUtil.subIndex(hit.docId(), context.searcher().getIndexReader().leaves());
-            LeafReaderContext readerContext = context.searcher().getIndexReader().leaves().get(readerIndex);
-            sourceLookup.setSegmentAndDocument(readerContext, hit.docId());
+        FieldValueRetriever fieldValueRetriever = fetchFieldsContext.fieldValueRetriever();
 
-            Set<String> ignoredFields = getIgnoredFields(hit);
-            Map<String, DocumentField> documentFields = fieldValueRetriever.retrieve(sourceLookup, ignoredFields);
-            for (Map.Entry<String, DocumentField> entry : documentFields.entrySet()) {
-                hit.setDocumentField(entry.getKey(), entry.getValue());
-            }
+        Set<String> ignoredFields = getIgnoredFields(hit);
+        Map<String, DocumentField> documentFields = fieldValueRetriever.retrieve(sourceLookup, ignoredFields);
+        for (Map.Entry<String, DocumentField> entry : documentFields.entrySet()) {
+            hit.setDocumentField(entry.getKey(), entry.getValue());
         }
     }