Add a simple 'fetch fields' phase. (#55639)

Currently the phase just looks up each field name in the _source and returns its
values in the 'fields' section of the response. There are several aspects that
need improvement -- this PR just lays out the initial class structure and tests.

Author: jtibshirani

rest-api-spec/src/main/resources/rest-api-spec/api/search.json

@@ -56,6 +56,10 @@
         "type":"string",
         "description":"The field to use as default where no field prefix is given in the query string"
       },
+      "fields": {
+        "type":"list",
+        "description":"A comma-separated list of fields to retrieve as part of each hit"
+      },
       "explain":{
         "type":"boolean",
         "description":"Specify whether to return detailed information about score computation as part of a hit"

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

@@ -0,0 +1,45 @@
+setup:
+  - skip:
+      version: " - 7.99.99"
+      reason: "fields retrieval is currently only implemented on master"
+  - do:
+      indices.create:
+          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
+
+  - do:
+      indices.refresh:
+          index: [ test ]
+
+---
+"Test basic field retrieval":
+  - do:
+      search:
+        index: test
+        body:
+          fields: [keyword, integer_range]
+
+  - is_true: hits.hits.0._id
+  - is_true: hits.hits.0._source
+
+  - match: { hits.hits.0.fields.keyword.0: first }
+  - match: { hits.hits.0.fields.keyword.1: second }
+
+  - match: { hits.hits.0.fields.integer_range.0.gte: 0 }
+  - match: { hits.hits.0.fields.integer_range.0.lte: 42 }

server/src/main/java/org/elasticsearch/action/search/SearchRequestBuilder.java

@@ -303,6 +303,11 @@ public SearchRequestBuilder addDocValueField(String name) {
         return addDocValueField(name, null);
     }
 
+    public SearchRequestBuilder addFetchField(String name) {
+        sourceBuilder().fetchField(name);
+        return this;
+    }
+
     /**
      * Adds a stored field to load and return (note, it must be stored) as part of the search request.
      */

server/src/main/java/org/elasticsearch/common/document/DocumentField.java

@@ -105,10 +105,9 @@ public void writeTo(StreamOutput out) throws IOException {
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startArray(name);
         for (Object value : values) {
-            // this call doesn't really need to support writing any kind of object.
-            // Stored fields values are converted using MappedFieldType#valueForDisplay.
-            // As a result they can either be Strings, Numbers, or Booleans, that's
-            // all.
+            // This call doesn't really need to support writing any kind of object, since the values
+            // here are always serializable to xContent. Each value could be a leaf types like a string,
+            // number, or boolean, a list of such values, or a map of such values with string keys.
             builder.value(value);
         }
         builder.endArray();

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

@@ -56,6 +56,7 @@
 import org.elasticsearch.search.fetch.FetchSearchResult;
 import org.elasticsearch.search.fetch.StoredFieldsContext;
 import org.elasticsearch.search.fetch.subphase.FetchDocValuesContext;
+import org.elasticsearch.search.fetch.subphase.FetchFieldsContext;
 import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
 import org.elasticsearch.search.fetch.subphase.ScriptFieldsContext;
 import org.elasticsearch.search.fetch.subphase.highlight.SearchContextHighlight;
@@ -111,6 +112,7 @@ final class DefaultSearchContext extends SearchContext {
     private ScriptFieldsContext scriptFields;
     private FetchSourceContext fetchSourceContext;
     private FetchDocValuesContext docValuesContext;
+    private FetchFieldsContext fetchFieldsContext;
     private int from = -1;
     private int size = -1;
     private SortAndFormats sort;
@@ -454,6 +456,17 @@ public SearchContext docValuesContext(FetchDocValuesContext docValuesContext) {
         return this;
     }
 
+    @Override
+    public FetchFieldsContext fetchFieldsContext() {
+        return fetchFieldsContext;
+    }
+
+    @Override
+    public SearchContext fetchFieldsContext(FetchFieldsContext fetchFieldsContext) {
+        this.fetchFieldsContext = fetchFieldsContext;
+        return this;
+    }
+
     @Override
     public ContextIndexSearcher searcher() {
         return this.searcher;

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

@@ -220,6 +220,7 @@
 import org.elasticsearch.search.fetch.FetchSubPhase;
 import org.elasticsearch.search.fetch.subphase.ExplainPhase;
 import org.elasticsearch.search.fetch.subphase.FetchDocValuesPhase;
+import org.elasticsearch.search.fetch.subphase.FetchFieldsPhase;
 import org.elasticsearch.search.fetch.subphase.FetchScorePhase;
 import org.elasticsearch.search.fetch.subphase.FetchSourcePhase;
 import org.elasticsearch.search.fetch.subphase.FetchVersionPhase;
@@ -731,6 +732,7 @@ private void registerFetchSubPhases(List<SearchPlugin> plugins) {
         registerFetchSubPhase(new FetchDocValuesPhase());
         registerFetchSubPhase(new ScriptFieldsPhase());
         registerFetchSubPhase(new FetchSourcePhase());
+        registerFetchSubPhase(new FetchFieldsPhase());
         registerFetchSubPhase(new FetchVersionPhase());
         registerFetchSubPhase(new SeqNoPrimaryTermPhase());
         registerFetchSubPhase(new MatchedQueriesPhase());

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

@@ -86,6 +86,7 @@
 import org.elasticsearch.search.fetch.ScrollQueryFetchSearchResult;
 import org.elasticsearch.search.fetch.ShardFetchRequest;
 import org.elasticsearch.search.fetch.subphase.FetchDocValuesContext;
+import org.elasticsearch.search.fetch.subphase.FetchFieldsContext;
 import org.elasticsearch.search.fetch.subphase.ScriptFieldsContext.ScriptField;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
 import org.elasticsearch.search.internal.AliasFilter;
@@ -918,6 +919,9 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
             FetchDocValuesContext docValuesContext = FetchDocValuesContext.create(context.mapperService(), source.docValueFields());
             context.docValuesContext(docValuesContext);
         }
+        if (source.fetchFields() != null) {
+            context.fetchFieldsContext(new FetchFieldsContext(source.fetchFields()));
+        }
         if (source.highlighter() != null) {
             HighlightBuilder highlightBuilder = source.highlighter();
             try {

server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java

@@ -20,6 +20,7 @@
 package org.elasticsearch.search.builder;
 
 import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.Version;
 import org.elasticsearch.common.Booleans;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.ParseField;
@@ -89,6 +90,7 @@ public final class SearchSourceBuilder implements Writeable, ToXContentObject, R
     public static final ParseField _SOURCE_FIELD = new ParseField("_source");
     public static final ParseField STORED_FIELDS_FIELD = new ParseField("stored_fields");
     public static final ParseField DOCVALUE_FIELDS_FIELD = new ParseField("docvalue_fields");
+    public static final ParseField FETCH_FIELDS_FIELD = new ParseField("fields");
     public static final ParseField SCRIPT_FIELDS_FIELD = new ParseField("script_fields");
     public static final ParseField SCRIPT_FIELD = new ParseField("script");
     public static final ParseField IGNORE_FAILURE_FIELD = new ParseField("ignore_failure");
@@ -165,6 +167,7 @@ public static HighlightBuilder highlight() {
     private List<FieldAndFormat> docValueFields;
     private List<ScriptField> scriptFields;
     private FetchSourceContext fetchSourceContext;
+    private List<String> fetchFields;
 
     private AggregatorFactories.Builder aggregations;
 
@@ -239,6 +242,10 @@ public SearchSourceBuilder(StreamInput in) throws IOException {
         sliceBuilder = in.readOptionalWriteable(SliceBuilder::new);
         collapse = in.readOptionalWriteable(CollapseBuilder::new);
         trackTotalHitsUpTo = in.readOptionalInt();
+
+        if (in.getVersion().onOrAfter(Version.V_8_0_0)) {
+            fetchFields = in.readOptionalStringList();
+        }
     }
 
     @Override
@@ -293,6 +300,10 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeOptionalWriteable(sliceBuilder);
         out.writeOptionalWriteable(collapse);
         out.writeOptionalInt(trackTotalHitsUpTo);
+
+        if (out.getVersion().onOrAfter(Version.V_8_0_0)) {
+            out.writeOptionalStringCollection(fetchFields);
+        }
     }
 
     /**
@@ -821,6 +832,24 @@ public SearchSourceBuilder docValueField(String name) {
         return docValueField(name, null);
     }
 
+    /**
+     * Gets the fields to load and return as part of the search request.
+     */
+    public List<String> fetchFields() {
+        return fetchFields;
+    }
+
+    /**
+     * Adds a field to load and return as part of the search request.
+     */
+    public SearchSourceBuilder fetchField(String fieldName) {
+        if (fetchFields == null) {
+            fetchFields = new ArrayList<>();
+        }
+        fetchFields.add(fieldName);
+        return this;
+    }
+
     /**
      * Adds a script field under the given name with the provided script.
      *
@@ -1116,6 +1145,11 @@ public void parseXContent(XContentParser parser, boolean checkTrailingTokens) th
                     while ((token = parser.nextToken()) != XContentParser.Token.END_ARRAY) {
                         docValueFields.add(FieldAndFormat.fromXContent(parser));
                     }
+                } else if (FETCH_FIELDS_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
+                    fetchFields = new ArrayList<>();
+                    while ((token = parser.nextToken()) != XContentParser.Token.END_ARRAY) {
+                        fetchFields.add(parser.text());
+                    }
                 } else if (INDICES_BOOST_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
                     while ((token = parser.nextToken()) != XContentParser.Token.END_ARRAY) {
                         indexBoosts.add(new IndexBoost(parser));
@@ -1223,6 +1257,10 @@ public XContentBuilder innerToXContent(XContentBuilder builder, Params params) t
             builder.endArray();
         }
 
+        if (fetchFields != null) {
+            builder.array(FETCH_FIELDS_FIELD.getPreferredName(), fetchFields);
+        }
+
         if (scriptFields != null) {
             builder.startObject(SCRIPT_FIELDS_FIELD.getPreferredName());
             for (ScriptField scriptField : scriptFields) {

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

@@ -0,0 +1,37 @@
+/*
+ * 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 java.util.List;
+
+/**
+ * The context needed to retrieve fields.
+ */
+public class FetchFieldsContext {
+
+    private final List<String> fields;
+
+    public FetchFieldsContext(List<String> fields) {
+        this.fields = fields;
+    }
+
+    public List<String> fields() {
+        return this.fields;
+    }
+}

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

@@ -0,0 +1,115 @@
+/*
+ * 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.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
+ * retrieves the field values from _source and returns them as document fields.
+ */
+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;
+        }
+
+        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() + "]");
+        }
+
+        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);
+        }
+
+        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.setDocumentField(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;
+    }
+
+    /**
+     * 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);
+            }
+        }
+        return result;
+    }
+}