Implement fields fetch for runtime fields

server/src/main/java/org/elasticsearch/index/mapper/DocValueFetcher.java

@@ -27,38 +27,21 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.function.Supplier;
 
 /**
  * Value fetcher that loads from doc values.
  */
 public final class DocValueFetcher implements ValueFetcher {
     private final DocValueFormat format;
-    private final Supplier<IndexFieldData<?>> ifdSource;
-    /**
-     * Field data implementation to load the values. Lazily initialized
-     * because it is expensive.
-     */
-    private IndexFieldData<?> ifd;
+    private final IndexFieldData<?> ifd;
     private Leaf leaf;
 
-    /**
-     * Build the fetcher.
-     * @param format the format to use
-     * @param ifdSource a supplier that will build the field data when we need it.
-     *                  We don't take the {@code IndexFieldData} directly here because
-     *                  it can be expensive to build and use a fair bit of memory so
-     *                  we wait until we're sure we need it.
-     */
-    public DocValueFetcher(DocValueFormat format, Supplier<IndexFieldData<?>> ifdSource) {
+    public DocValueFetcher(DocValueFormat format, IndexFieldData<?> ifd) {
         this.format = format;
-        this.ifdSource = ifdSource;
+        this.ifd = ifd;
     }
 
     public void setNextReader(LeafReaderContext context) {
-        if (ifd == null) {
-            ifd = ifdSource.get();
-        }
         leaf = ifd.load(context).getLeafValueFetcher(format);
     }
 

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

@@ -62,7 +62,7 @@ public FetchSubPhaseProcessor getProcessor(SearchContext context) throws IOExcep
             if (ft == null) {
                 continue;
             }
-            ValueFetcher fetcher = new DocValueFetcher(ft.docValueFormat(fieldAndFormat.format, null), () -> context.getForField(ft));
+            ValueFetcher fetcher = new DocValueFetcher(ft.docValueFormat(fieldAndFormat.format, null), context.getForField(ft));
             fields.add(new DocValueField(fieldAndFormat.field, fetcher));
         }
 

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

@@ -23,6 +23,7 @@
 import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.util.List;
+import java.util.function.BiFunction;
 
 /**
  * The context needed to retrieve fields.
@@ -36,21 +37,21 @@ public static FetchFieldsContext create(
     ) {
         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 + "]");
+            throw new IllegalArgumentException(
+                "Unable to retrieve the requested [fields] since _source is " + "disabled in the mappings for index [" + indexName + "]"
+            );
         }
 
-        FieldValueRetriever fieldValueRetriever = FieldValueRetriever.create(mapperService, searchLookup, fields);
-        return new FetchFieldsContext(fieldValueRetriever);
+        return new FetchFieldsContext((m, s) -> FieldValueRetriever.create(m, s, fields));
     }
 
-    private final FieldValueRetriever fieldValueRetriever;
+    private final BiFunction<MapperService, SearchLookup, FieldValueRetriever> fieldValueRetriever;
 
-    private FetchFieldsContext(FieldValueRetriever fieldValueRetriever) {
+    private FetchFieldsContext(BiFunction<MapperService, SearchLookup, FieldValueRetriever> fieldValueRetriever) {
         this.fieldValueRetriever = fieldValueRetriever;
     }
 
-    public FieldValueRetriever fieldValueRetriever() {
-        return fieldValueRetriever;
+    public FieldValueRetriever fieldValueRetriever(MapperService mapperService, SearchLookup searchLookup) {
+        return fieldValueRetriever.apply(mapperService, searchLookup);
     }
 }

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

@@ -45,20 +45,23 @@ public FetchSubPhaseProcessor getProcessor(SearchContext searchContext) {
         if (fetchFieldsContext == null) {
             return null;
         }
+        FieldValueRetriever retriever = fetchFieldsContext.fieldValueRetriever(
+            searchContext.mapperService(),
+            searchContext.getQueryShardContext().fetchLookup()
+        );
         return new FetchSubPhaseProcessor() {
             @Override
             public void setNextReader(LeafReaderContext readerContext) {
-                fetchFieldsContext.fieldValueRetriever().setNextReader(readerContext);
+                retriever.setNextReader(readerContext);
             }
 
             @Override
             public void process(HitContext hitContext) throws IOException {
                 SearchHit hit = hitContext.hit();
                 SourceLookup sourceLookup = hitContext.sourceLookup();
-                FieldValueRetriever fieldValueRetriever = fetchFieldsContext.fieldValueRetriever();
 
                 Set<String> ignoredFields = getIgnoredFields(hit);
-                Map<String, DocumentField> documentFields = fieldValueRetriever.retrieve(sourceLookup, ignoredFields);
+                Map<String, DocumentField> documentFields = retriever.retrieve(sourceLookup, ignoredFields);
                 for (Map.Entry<String, DocumentField> entry : documentFields.entrySet()) {
                     hit.setDocumentField(entry.getKey(), entry.getValue());
                 }

test/framework/src/main/java/org/elasticsearch/index/mapper/MapperTestCase.java

@@ -192,7 +192,7 @@ protected final List<?> fetchFromDocValues(FieldMapper mapper, DocValueFormat fo
         }, iw -> {
             IndexSearcher indexSearcher = newSearcher(iw);
             SearchLookup lookup = new SearchLookup(mapperService, fieldDataLookup);
-            ValueFetcher valueFetcher = new DocValueFetcher(format, () -> lookup.doc().getForField(mapper.fieldType()));
+            ValueFetcher valueFetcher = new DocValueFetcher(format, lookup.doc().getForField(mapper.fieldType()));
             indexSearcher.search(new MatchAllDocsQuery(), new Collector() {
                 @Override
                 public ScoreMode scoreMode() {

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/mapper/RuntimeScriptFieldMapper.java

@@ -78,7 +78,7 @@ protected void parseCreateField(ParseContext context) {
 
     @Override
     public ValueFetcher valueFetcher(MapperService mapperService, SearchLookup lookup, String format) {
-        return new DocValueFetcher(fieldType().docValueFormat(format, null), () -> lookup.doc().getForField(fieldType()));
+        return new DocValueFetcher(fieldType().docValueFormat(format, null), lookup.doc().getForField(fieldType()));
     }
 
     @Override