From fc2521e60ea77e732d227fac76cd7e8e7627e445 Mon Sep 17 00:00:00 2001
From: Julie Tibshirani <julie.tibshirani@elastic.co>
Date: Tue, 16 Jun 2020 14:36:44 -0700
Subject: [PATCH] Consolidate logic around doc values loading.

This improves modularity and also fixes some issues when `docvalues_fields` is
used within `inner_hits` or the `top_hits` agg:
* We previously didn't resolve wildcards in field names.
* We also forgot to enforce the limit `index.max_docvalue_fields_search`.
---
 .../search/fetch/subphase/InnerHitsIT.java    |  2 +-
 .../index/query/InnerHitContextBuilder.java   |  4 ++-
 .../elasticsearch/search/SearchService.java   | 19 ++------------
 .../metrics/TopHitsAggregatorFactory.java     |  3 ++-
 .../fetch/subphase/FetchDocValuesContext.java | 26 ++++++++++++++++++-
 5 files changed, 33 insertions(+), 21 deletions(-)

diff --git a/server/src/internalClusterTest/java/org/elasticsearch/search/fetch/subphase/InnerHitsIT.java b/server/src/internalClusterTest/java/org/elasticsearch/search/fetch/subphase/InnerHitsIT.java
index 93fbec203fe0a..16dbc6f93bfd1 100644
--- a/server/src/internalClusterTest/java/org/elasticsearch/search/fetch/subphase/InnerHitsIT.java
+++ b/server/src/internalClusterTest/java/org/elasticsearch/search/fetch/subphase/InnerHitsIT.java
@@ -165,7 +165,7 @@ public void testSimpleNested() throws Exception {
                 .setQuery(nestedQuery("comments", matchQuery("comments.message", "fox"), ScoreMode.Avg).innerHit(
                         new InnerHitBuilder().setHighlightBuilder(new HighlightBuilder().field("comments.message"))
                                 .setExplain(true)
-                                .addDocValueField("comments.message")
+                                .addDocValueField("comments.mes*")
                                 .addScriptField("script",
                                         new Script(ScriptType.INLINE, MockScriptEngine.NAME, "5", Collections.emptyMap()))
                                 .setSize(1))).get();
diff --git a/server/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java b/server/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java
index acf55d6ca4330..717ace6bb045b 100644
--- a/server/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java
+++ b/server/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java
@@ -88,7 +88,9 @@ protected void setupInnerHitsContext(QueryShardContext queryShardContext,
             innerHitsContext.storedFieldsContext(innerHitBuilder.getStoredFieldsContext());
         }
         if (innerHitBuilder.getDocValueFields() != null) {
-            innerHitsContext.docValuesContext(new FetchDocValuesContext(innerHitBuilder.getDocValueFields()));
+            FetchDocValuesContext docValuesContext = FetchDocValuesContext.create(
+                queryShardContext.getMapperService(), innerHitBuilder.getDocValueFields());
+            innerHitsContext.docValuesContext(docValuesContext);
         }
         if (innerHitBuilder.getScriptFields() != null) {
             for (SearchSourceBuilder.ScriptField field : innerHitBuilder.getScriptFields()) {
diff --git a/server/src/main/java/org/elasticsearch/search/SearchService.java b/server/src/main/java/org/elasticsearch/search/SearchService.java
index 0f3ef7cd4b79b..2b8dab733bf0d 100644
--- a/server/src/main/java/org/elasticsearch/search/SearchService.java
+++ b/server/src/main/java/org/elasticsearch/search/SearchService.java
@@ -114,8 +114,6 @@
 import org.elasticsearch.transport.TransportRequest;
 
 import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -917,21 +915,8 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
             context.fetchSourceContext(source.fetchSource());
         }
         if (source.docValueFields() != null) {
-            List<FetchDocValuesContext.FieldAndFormat> docValueFields = new ArrayList<>();
-            for (FetchDocValuesContext.FieldAndFormat format : source.docValueFields()) {
-                Collection<String> fieldNames = context.mapperService().simpleMatchToFullName(format.field);
-                for (String fieldName: fieldNames) {
-                   docValueFields.add(new FetchDocValuesContext.FieldAndFormat(fieldName, format.format));
-                }
-            }
-            int maxAllowedDocvalueFields = context.mapperService().getIndexSettings().getMaxDocvalueFields();
-            if (docValueFields.size() > maxAllowedDocvalueFields) {
-                throw new IllegalArgumentException(
-                    "Trying to retrieve too many docvalue_fields. Must be less than or equal to: [" + maxAllowedDocvalueFields
-                        + "] but was [" + docValueFields.size() + "]. This limit can be set by changing the ["
-                        + IndexSettings.MAX_DOCVALUE_FIELDS_SEARCH_SETTING.getKey() + "] index level setting.");
-            }
-            context.docValuesContext(new FetchDocValuesContext(docValueFields));
+            FetchDocValuesContext docValuesContext = FetchDocValuesContext.create(context.mapperService(), source.docValueFields());
+            context.docValuesContext(docValuesContext);
         }
         if (source.highlighter() != null) {
             HighlightBuilder highlightBuilder = source.highlighter();
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/TopHitsAggregatorFactory.java b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/TopHitsAggregatorFactory.java
index 142141b21b6a5..a0a480fbc55c3 100644
--- a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/TopHitsAggregatorFactory.java
+++ b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/TopHitsAggregatorFactory.java
@@ -106,7 +106,8 @@ public Aggregator createInternal(SearchContext searchContext,
             subSearchContext.storedFieldsContext(storedFieldsContext);
         }
         if (docValueFields != null) {
-            subSearchContext.docValuesContext(new FetchDocValuesContext(docValueFields));
+            FetchDocValuesContext docValuesContext = FetchDocValuesContext.create(searchContext.mapperService(), docValueFields);
+            subSearchContext.docValuesContext(docValuesContext);
         }
         for (ScriptFieldsContext.ScriptField field : scriptFields) {
             subSearchContext.scriptFields().add(field);
diff --git a/server/src/main/java/org/elasticsearch/search/fetch/subphase/FetchDocValuesContext.java b/server/src/main/java/org/elasticsearch/search/fetch/subphase/FetchDocValuesContext.java
index 5eaedac992fb1..c4449d0137779 100644
--- a/server/src/main/java/org/elasticsearch/search/fetch/subphase/FetchDocValuesContext.java
+++ b/server/src/main/java/org/elasticsearch/search/fetch/subphase/FetchDocValuesContext.java
@@ -27,8 +27,12 @@
 import org.elasticsearch.common.xcontent.XContent;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.common.xcontent.XContentParser.Token;
+import org.elasticsearch.index.IndexSettings;
+import org.elasticsearch.index.mapper.MapperService;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import java.util.Objects;
 
@@ -106,7 +110,27 @@ public boolean equals(Object obj) {
 
     private final List<FieldAndFormat> fields;
 
-    public FetchDocValuesContext(List<FieldAndFormat> fields) {
+    public static FetchDocValuesContext create(MapperService mapperService,
+                                               List<FieldAndFormat> fieldPatterns) {
+        List<FieldAndFormat> fields = new ArrayList<>();
+        for (FieldAndFormat field : fieldPatterns) {
+            Collection<String> fieldNames = mapperService.simpleMatchToFullName(field.field);
+            for (String fieldName: fieldNames) {
+                fields.add(new FieldAndFormat(fieldName, field.format));
+            }
+        }
+        int maxAllowedDocvalueFields = mapperService.getIndexSettings().getMaxDocvalueFields();
+        if (fields.size() > maxAllowedDocvalueFields) {
+            throw new IllegalArgumentException(
+                "Trying to retrieve too many docvalue_fields. Must be less than or equal to: [" + maxAllowedDocvalueFields
+                    + "] but was [" + fields.size() + "]. This limit can be set by changing the ["
+                    + IndexSettings.MAX_DOCVALUE_FIELDS_SEARCH_SETTING.getKey() + "] index level setting.");
+        }
+
+        return new FetchDocValuesContext(fields);
+    }
+
+    FetchDocValuesContext(List<FieldAndFormat> fields) {
         this.fields = fields;
     }