Make the fetch fields phase easier to test. (#55756)

This commit pulls out a FieldValueRetriever object, which retrieves specific
fields given a document's source. The new object makes it easier to unit test
the logic, and will help keep FetchFieldsPhase from growing too complex as we
add more functionality.

Author: jtibshirani

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

@@ -2,33 +2,34 @@ setup:
   - skip:
       version: " - 7.99.99"
       reason: "fields retrieval is currently only implemented on master"
+
+---
+"Test basic field retrieval":
   - do:
       indices.create:
-          index:  test
-          body:
-            mappings:
-              properties:
-                keyword:
-                  type: keyword
-                integer_range:
-                  type: integer_range
+        index:  test
+        body:
+          mappings:
+            properties:
+              keyword:
+                type: keyword
+              integer_range:
+                type: integer_range
 
   - do:
       index:
-          index:  test
-          id:     1
-          body:
-            keyword: [ "first", "second" ]
-            integer_range:
-              gte: 0
-              lte: 42
+        index:  test
+        id:     1
+        body:
+          keyword: [ "first", "second" ]
+          integer_range:
+            gte: 0
+            lte: 42
 
   - do:
       indices.refresh:
-          index: [ test ]
+        index: [ test ]
 
----
-"Test basic field retrieval":
   - do:
       search:
         index: test
@@ -43,3 +44,35 @@ setup:
 
   - match: { hits.hits.0.fields.integer_range.0.gte: 0 }
   - match: { hits.hits.0.fields.integer_range.0.lte: 42 }
+
+---
+"Test disable source":
+  - do:
+      indices.create:
+        index:  test
+        body:
+          settings:
+            number_of_shards: 1
+          mappings:
+            _source:
+              enabled: false
+            properties:
+              keyword:
+                type: keyword
+
+  - do:
+      index:
+        index:  test
+        id:     1
+        body:
+          keyword: [ "value" ]
+
+  - do:
+      catch: bad_request
+      search:
+        index: test
+        body:
+          fields: [keyword]
+  - match: { error.root_cause.0.type: "illegal_argument_exception" }
+  - match: { error.root_cause.0.reason: "Unable to retrieve the requested [fields] since _source is disabled
+        in the mappings for index [test]" }

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

@@ -22,20 +22,13 @@
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.ReaderUtil;
 import org.elasticsearch.common.document.DocumentField;
-import org.elasticsearch.common.xcontent.support.XContentMapValues;
 import org.elasticsearch.index.mapper.DocumentMapper;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.fetch.FetchSubPhase;
 import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.lookup.SourceLookup;
 
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.List;
 import java.util.Map;
-import java.util.Set;
-import java.util.function.Function;
 
 /**
  * A fetch sub-phase for high-level field retrieval. Given a list of fields, it
@@ -45,14 +38,6 @@ public final class FetchFieldsPhase implements FetchSubPhase {
 
     @Override
     public void hitsExecute(SearchContext context, SearchHit[] hits) {
-        hitsExecute(context, hit -> getSourceLookup(context, hit), hits);
-    }
-
-    // Visible for testing.
-    @SuppressWarnings("unchecked")
-    void hitsExecute(SearchContext context,
-                     Function<SearchHit, SourceLookup> sourceProvider,
-                     SearchHit[] hits) {
         FetchFieldsContext fetchFieldsContext = context.fetchFieldsContext();
         if (fetchFieldsContext == null || fetchFieldsContext.fields().isEmpty()) {
             return;
@@ -64,52 +49,18 @@ void hitsExecute(SearchContext context,
                 "disabled in the mappings for index [" + context.indexShard().shardId().getIndexName() + "]");
         }
 
-        Set<String> fields = new HashSet<>();
-        for (String fieldPattern : context.fetchFieldsContext().fields()) {
-            if (documentMapper.objectMappers().containsKey(fieldPattern)) {
-                continue;
-            }
-            Collection<String> concreteFields = context.mapperService().simpleMatchToFullName(fieldPattern);
-            fields.addAll(concreteFields);
-        }
+        SourceLookup sourceLookup = context.lookup().source();
+        FieldValueRetriever fieldValueRetriever = FieldValueRetriever.create(
+            context.mapperService(),
+            fetchFieldsContext.fields());
 
         for (SearchHit hit : hits) {
-            SourceLookup sourceLookup = sourceProvider.apply(hit);
-            Map<String, Object> valuesByField = extractValues(sourceLookup, fields);
-
-            for (Map.Entry<String, Object> entry : valuesByField.entrySet()) {
-                String field = entry.getKey();
-                Object value = entry.getValue();
-                List<Object> values = value instanceof List
-                    ? (List<Object>) value
-                    : List.of(value);
-
-                DocumentField documentField = new DocumentField(field, values);
-                hit.setField(field, documentField);
-            }
-        }
-    }
-
-    private SourceLookup getSourceLookup(SearchContext context, SearchHit hit) {
-        SourceLookup sourceLookup = context.lookup().source();
-        int readerIndex = ReaderUtil.subIndex(hit.docId(), context.searcher().getIndexReader().leaves());
-        LeafReaderContext readerContext = context.searcher().getIndexReader().leaves().get(readerIndex);
-        sourceLookup.setSegmentAndDocument(readerContext, hit.docId());
-        return sourceLookup;
-    }
+            int readerIndex = ReaderUtil.subIndex(hit.docId(), context.searcher().getIndexReader().leaves());
+            LeafReaderContext readerContext = context.searcher().getIndexReader().leaves().get(readerIndex);
+            sourceLookup.setSegmentAndDocument(readerContext, hit.docId());
 
-    /**
-     * For each of the provided paths, return its value in the source. Note that in contrast with
-     * {@link SourceLookup#extractRawValues}, array and object values can be returned.
-     */
-    private Map<String, Object> extractValues(SourceLookup sourceLookup, Collection<String> paths) {
-        Map<String, Object> result = new HashMap<>(paths.size());
-        for (String path : paths) {
-            Object value = XContentMapValues.extractValue(path, sourceLookup);
-            if (value != null) {
-                result.put(path, value);
-            }
+            Map<String, DocumentField> fieldValues = fieldValueRetriever.retrieve(sourceLookup);
+            hit.fields(fieldValues);
         }
-        return result;
     }
 }

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

@@ -0,0 +1,87 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.search.fetch.subphase;
+
+import org.elasticsearch.common.document.DocumentField;
+import org.elasticsearch.common.xcontent.support.XContentMapValues;
+import org.elasticsearch.index.mapper.DocumentMapper;
+import org.elasticsearch.index.mapper.MapperService;
+import org.elasticsearch.search.lookup.SourceLookup;
+
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public class FieldValueRetriever {
+    private final Set<String> fields;
+
+    public static FieldValueRetriever create(MapperService mapperService,
+                                             Collection<String> fieldPatterns) {
+        Set<String> fields = new HashSet<>();
+        DocumentMapper documentMapper = mapperService.documentMapper();
+
+        for (String fieldPattern : fieldPatterns) {
+            if (documentMapper.objectMappers().containsKey(fieldPattern)) {
+                continue;
+            }
+            Collection<String> concreteFields = mapperService.simpleMatchToFullName(fieldPattern);
+            fields.addAll(concreteFields);
+        }
+        return new FieldValueRetriever(fields);
+    }
+
+    private FieldValueRetriever(Set<String> fields) {
+        this.fields = fields;
+    }
+
+    @SuppressWarnings("unchecked")
+    public Map<String, DocumentField> retrieve(SourceLookup sourceLookup) {
+        Map<String, DocumentField> result = new HashMap<>();
+        Map<String, Object> sourceValues = extractValues(sourceLookup, this.fields);
+
+        for (Map.Entry<String, Object> entry : sourceValues.entrySet()) {
+            String field = entry.getKey();
+            Object value = entry.getValue();
+            List<Object> values = value instanceof List
+                ? (List<Object>) value
+                : List.of(value);
+            result.put(field, new DocumentField(field, values));
+        }
+        return result;
+    }
+
+    /**
+     * For each of the provided paths, return its value in the source. Note that in contrast with
+     * {@link SourceLookup#extractRawValues}, array and object values can be returned.
+     */
+    private static Map<String, Object> extractValues(SourceLookup sourceLookup, Collection<String> paths) {
+        Map<String, Object> result = new HashMap<>(paths.size());
+        for (String path : paths) {
+            Object value = XContentMapValues.extractValue(path, sourceLookup);
+            if (value != null) {
+                result.put(path, value);
+            }
+        }
+        return result;
+    }
+}