make docValueFields similar to other standard sub fetch phases

Given that doc value fields is our own fetch sub phase, it doesn't need to be implemented like if it was plugged in from the outside. It doesn't need its own fetch sub phase context, but it can just be an instance member in SearchContext

8 files changed, 50 insertions, 69 deletions

diff --git a/core/src/main/java/org/elasticsearch/index/query/InnerHitBuilder.java b/core/src/main/java/org/elasticsearch/index/query/InnerHitBuilder.java
index 27b0138e1e..4516dfde69 100644
--- a/core/src/main/java/org/elasticsearch/index/query/InnerHitBuilder.java
+++ b/core/src/main/java/org/elasticsearch/index/query/InnerHitBuilder.java
@@ -36,7 +36,6 @@ import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.builder.SearchSourceBuilder.ScriptField;
 import org.elasticsearch.search.fetch.StoredFieldsContext;
 import org.elasticsearch.search.fetch.subphase.DocValueFieldsContext;
-import org.elasticsearch.search.fetch.subphase.DocValueFieldsFetchSubPhase;
 import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
 import org.elasticsearch.search.fetch.subphase.InnerHitsContext;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
@@ -572,12 +571,7 @@ public final class InnerHitBuilder extends ToXContentToBytes implements Writeabl
             innerHitsContext.storedFieldsContext(storedFieldsContext);
         }
         if (docValueFields != null) {
-            DocValueFieldsContext docValueFieldsContext = innerHitsContext
-                    .getFetchSubPhaseContext(DocValueFieldsFetchSubPhase.CONTEXT_FACTORY);
-            for (String field : docValueFields) {
-                docValueFieldsContext.add(new DocValueFieldsContext.DocValueField(field));
-            }
-            docValueFieldsContext.setHitExecutionNeeded(true);
+            innerHitsContext.docValueFieldsContext(new DocValueFieldsContext(docValueFields));
         }
         if (scriptFields != null) {
             for (ScriptField field : scriptFields) {
diff --git a/core/src/main/java/org/elasticsearch/search/SearchService.java b/core/src/main/java/org/elasticsearch/search/SearchService.java
index daebe480e0..223595cc13 100644
--- a/core/src/main/java/org/elasticsearch/search/SearchService.java
+++ b/core/src/main/java/org/elasticsearch/search/SearchService.java
@@ -67,8 +67,6 @@ import org.elasticsearch.search.fetch.QueryFetchSearchResult;
 import org.elasticsearch.search.fetch.ScrollQueryFetchSearchResult;
 import org.elasticsearch.search.fetch.ShardFetchRequest;
 import org.elasticsearch.search.fetch.subphase.DocValueFieldsContext;
-import org.elasticsearch.search.fetch.subphase.DocValueFieldsContext.DocValueField;
-import org.elasticsearch.search.fetch.subphase.DocValueFieldsFetchSubPhase;
 import org.elasticsearch.search.fetch.subphase.ScriptFieldsContext.ScriptField;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
 import org.elasticsearch.search.internal.DefaultSearchContext;
@@ -732,11 +730,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
             context.fetchSourceContext(source.fetchSource());
         }
         if (source.docValueFields() != null) {
-            DocValueFieldsContext docValuesFieldsContext = context.getFetchSubPhaseContext(DocValueFieldsFetchSubPhase.CONTEXT_FACTORY);
-            for (String field : source.docValueFields()) {
-                docValuesFieldsContext.add(new DocValueField(field));
-            }
-            docValuesFieldsContext.setHitExecutionNeeded(true);
+            context.docValueFieldsContext(new DocValueFieldsContext(source.docValueFields()));
         }
         if (source.highlighter() != null) {
             HighlightBuilder highlightBuilder = source.highlighter();
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregatorFactory.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregatorFactory.java
index b3389322d9..7c6a743a20 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregatorFactory.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregatorFactory.java
@@ -19,7 +19,6 @@
 
 package org.elasticsearch.search.aggregations.metrics.tophits;
 
-import org.elasticsearch.search.fetch.StoredFieldsContext;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.SearchScript;
 import org.elasticsearch.search.aggregations.Aggregator;
@@ -29,9 +28,8 @@ import org.elasticsearch.search.aggregations.InternalAggregation.Type;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
 import org.elasticsearch.search.aggregations.support.AggregationContext;
 import org.elasticsearch.search.builder.SearchSourceBuilder.ScriptField;
+import org.elasticsearch.search.fetch.StoredFieldsContext;
 import org.elasticsearch.search.fetch.subphase.DocValueFieldsContext;
-import org.elasticsearch.search.fetch.subphase.DocValueFieldsContext.DocValueField;
-import org.elasticsearch.search.fetch.subphase.DocValueFieldsFetchSubPhase;
 import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
 import org.elasticsearch.search.internal.SubSearchContext;
@@ -98,12 +96,7 @@ public class TopHitsAggregatorFactory extends AggregatorFactory<TopHitsAggregato
             subSearchContext.storedFieldsContext(storedFieldsContext);
         }
         if (docValueFields != null) {
-            DocValueFieldsContext docValueFieldsContext = subSearchContext
-                    .getFetchSubPhaseContext(DocValueFieldsFetchSubPhase.CONTEXT_FACTORY);
-            for (String field : docValueFields) {
-                docValueFieldsContext.add(new DocValueField(field));
-            }
-            docValueFieldsContext.setHitExecutionNeeded(true);
+            subSearchContext.docValueFieldsContext(new DocValueFieldsContext(docValueFields));
         }
         if (scriptFields != null) {
             for (ScriptField field : scriptFields) {
diff --git a/core/src/main/java/org/elasticsearch/search/fetch/subphase/DocValueFieldsContext.java b/core/src/main/java/org/elasticsearch/search/fetch/subphase/DocValueFieldsContext.java
index 54185734f9..325d28e459 100644
--- a/core/src/main/java/org/elasticsearch/search/fetch/subphase/DocValueFieldsContext.java
+++ b/core/src/main/java/org/elasticsearch/search/fetch/subphase/DocValueFieldsContext.java
@@ -18,38 +18,23 @@
  */
 package org.elasticsearch.search.fetch.subphase;
 
-import org.elasticsearch.search.fetch.FetchSubPhaseContext;
-
-import java.util.ArrayList;
 import java.util.List;
 
 /**
  * All the required context to pull a field from the doc values.
  */
-public class DocValueFieldsContext extends FetchSubPhaseContext {
-
-    public static class DocValueField {
-        private final String name;
+public class DocValueFieldsContext {
 
-        public DocValueField(String name) {
-            this.name = name;
-        }
-
-        public String name() {
-            return name;
-        }
-    }
-
-    private List<DocValueField> fields = new ArrayList<>();
-
-    public DocValueFieldsContext() {
-    }
+    private final List<String> fields;
 
-    public void add(DocValueField field) {
-        this.fields.add(field);
+    public DocValueFieldsContext(List<String> fields) {
+        this.fields = fields;
     }
 
-    public List<DocValueField> fields() {
+    /**
+     * Returns the required docvalue fields
+     */
+    public List<String> fields() {
         return this.fields;
     }
 }
diff --git a/core/src/main/java/org/elasticsearch/search/fetch/subphase/DocValueFieldsFetchSubPhase.java b/core/src/main/java/org/elasticsearch/search/fetch/subphase/DocValueFieldsFetchSubPhase.java
index 803cbb4348..befce94a9e 100644
--- a/core/src/main/java/org/elasticsearch/search/fetch/subphase/DocValueFieldsFetchSubPhase.java
+++ b/core/src/main/java/org/elasticsearch/search/fetch/subphase/DocValueFieldsFetchSubPhase.java
@@ -36,35 +36,21 @@ import java.util.HashMap;
  */
 public final class DocValueFieldsFetchSubPhase implements FetchSubPhase {
 
-    public static final String NAME = "docvalue_fields";
-    public static final ContextFactory<DocValueFieldsContext> CONTEXT_FACTORY = new ContextFactory<DocValueFieldsContext>() {
-
-        @Override
-        public String getName() {
-            return NAME;
-        }
-
-        @Override
-        public DocValueFieldsContext newContextInstance() {
-            return new DocValueFieldsContext();
-        }
-    };
-
     @Override
     public void hitExecute(SearchContext context, HitContext hitContext) {
-        if (context.getFetchSubPhaseContext(CONTEXT_FACTORY).hitExecutionNeeded() == false) {
+        if (context.docValueFieldsContext() == null) {
             return;
         }
-        for (DocValueFieldsContext.DocValueField field : context.getFetchSubPhaseContext(CONTEXT_FACTORY).fields()) {
+        for (String field : context.docValueFieldsContext().fields()) {
             if (hitContext.hit().fieldsOrNull() == null) {
                 hitContext.hit().fields(new HashMap<>(2));
             }
-            SearchHitField hitField = hitContext.hit().fields().get(field.name());
+            SearchHitField hitField = hitContext.hit().fields().get(field);
             if (hitField == null) {
-                hitField = new InternalSearchHitField(field.name(), new ArrayList<>(2));
-                hitContext.hit().fields().put(field.name(), hitField);
+                hitField = new InternalSearchHitField(field, new ArrayList<>(2));
+                hitContext.hit().fields().put(field, hitField);
             }
-            MappedFieldType fieldType = context.mapperService().fullName(field.name());
+            MappedFieldType fieldType = context.mapperService().fullName(field);
             if (fieldType != null) {
                 AtomicFieldData data = context.fieldData().getForField(fieldType).load(hitContext.readerContext());
                 ScriptDocValues values = data.getScriptValues();
diff --git a/core/src/main/java/org/elasticsearch/search/internal/DefaultSearchContext.java b/core/src/main/java/org/elasticsearch/search/internal/DefaultSearchContext.java
index 780352508b..edce266c95 100644
--- a/core/src/main/java/org/elasticsearch/search/internal/DefaultSearchContext.java
+++ b/core/src/main/java/org/elasticsearch/search/internal/DefaultSearchContext.java
@@ -61,6 +61,7 @@ import org.elasticsearch.search.fetch.FetchPhase;
 import org.elasticsearch.search.fetch.FetchSearchResult;
 import org.elasticsearch.search.fetch.FetchSubPhase;
 import org.elasticsearch.search.fetch.FetchSubPhaseContext;
+import org.elasticsearch.search.fetch.subphase.DocValueFieldsContext;
 import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
 import org.elasticsearch.search.fetch.subphase.ScriptFieldsContext;
 import org.elasticsearch.search.fetch.subphase.highlight.SearchContextHighlight;
@@ -110,6 +111,7 @@ public class DefaultSearchContext extends SearchContext {
     private StoredFieldsContext storedFields;
     private ScriptFieldsContext scriptFields;
     private FetchSourceContext fetchSourceContext;
+    private DocValueFieldsContext docValueFieldsContext;
     private int from = -1;
     private int size = -1;
     private SortAndFormats sort;
@@ -471,6 +473,17 @@ public class DefaultSearchContext extends SearchContext {
     }
 
     @Override
+    public DocValueFieldsContext docValueFieldsContext() {
+        return docValueFieldsContext;
+    }
+
+    @Override
+    public SearchContext docValueFieldsContext(DocValueFieldsContext docValueFieldsContext) {
+        this.docValueFieldsContext = docValueFieldsContext;
+        return this;
+    }
+
+    @Override
     public ContextIndexSearcher searcher() {
         return this.searcher;
     }
diff --git a/core/src/main/java/org/elasticsearch/search/internal/SearchContext.java b/core/src/main/java/org/elasticsearch/search/internal/SearchContext.java
index 65fe7ddad1..543f97e638 100644
--- a/core/src/main/java/org/elasticsearch/search/internal/SearchContext.java
+++ b/core/src/main/java/org/elasticsearch/search/internal/SearchContext.java
@@ -22,9 +22,7 @@ package org.elasticsearch.search.internal;
 import org.apache.lucene.search.Collector;
 import org.apache.lucene.search.FieldDoc;
 import org.apache.lucene.search.Query;
-import org.apache.lucene.store.AlreadyClosedException;
 import org.apache.lucene.util.Counter;
-import org.apache.lucene.util.RefCount;
 import org.elasticsearch.action.search.SearchType;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.ParseFieldMatcher;
@@ -43,7 +41,6 @@ import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.ObjectMapper;
 import org.elasticsearch.index.query.ParsedQuery;
 import org.elasticsearch.index.query.QueryShardContext;
-import org.elasticsearch.search.fetch.StoredFieldsContext;
 import org.elasticsearch.index.shard.IndexShard;
 import org.elasticsearch.index.similarity.SimilarityService;
 import org.elasticsearch.script.ScriptService;
@@ -54,6 +51,8 @@ import org.elasticsearch.search.fetch.FetchPhase;
 import org.elasticsearch.search.fetch.FetchSearchResult;
 import org.elasticsearch.search.fetch.FetchSubPhase;
 import org.elasticsearch.search.fetch.FetchSubPhaseContext;
+import org.elasticsearch.search.fetch.StoredFieldsContext;
+import org.elasticsearch.search.fetch.subphase.DocValueFieldsContext;
 import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
 import org.elasticsearch.search.fetch.subphase.InnerHitsContext;
 import org.elasticsearch.search.fetch.subphase.ScriptFieldsContext;
@@ -226,6 +225,10 @@ public abstract class SearchContext extends AbstractRefCounted implements Releas
 
     public abstract SearchContext fetchSourceContext(FetchSourceContext fetchSourceContext);
 
+    public abstract DocValueFieldsContext docValueFieldsContext();
+
+    public abstract SearchContext docValueFieldsContext(DocValueFieldsContext docValueFieldsContext);
+
     public abstract ContextIndexSearcher searcher();
 
     public abstract IndexShard indexShard();
diff --git a/core/src/main/java/org/elasticsearch/search/internal/SubSearchContext.java b/core/src/main/java/org/elasticsearch/search/internal/SubSearchContext.java
index 077a0941a0..8ba76971de 100644
--- a/core/src/main/java/org/elasticsearch/search/internal/SubSearchContext.java
+++ b/core/src/main/java/org/elasticsearch/search/internal/SubSearchContext.java
@@ -25,6 +25,7 @@ import org.elasticsearch.index.query.ParsedQuery;
 import org.elasticsearch.search.fetch.StoredFieldsContext;
 import org.elasticsearch.search.aggregations.SearchContextAggregations;
 import org.elasticsearch.search.fetch.FetchSearchResult;
+import org.elasticsearch.search.fetch.subphase.DocValueFieldsContext;
 import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
 import org.elasticsearch.search.fetch.subphase.ScriptFieldsContext;
 import org.elasticsearch.search.fetch.subphase.highlight.SearchContextHighlight;
@@ -60,6 +61,7 @@ public class SubSearchContext extends FilteredSearchContext {
     private StoredFieldsContext storedFields;
     private ScriptFieldsContext scriptFields;
     private FetchSourceContext fetchSourceContext;
+    private DocValueFieldsContext docValueFieldsContext;
     private SearchContextHighlight highlight;
 
     private boolean explain;
@@ -155,6 +157,17 @@ public class SubSearchContext extends FilteredSearchContext {
     }
 
     @Override
+    public DocValueFieldsContext docValueFieldsContext() {
+        return docValueFieldsContext;
+    }
+
+    @Override
+    public SearchContext docValueFieldsContext(DocValueFieldsContext docValueFieldsContext) {
+        this.docValueFieldsContext = docValueFieldsContext;
+        return this;
+    }
+
+    @Override
     public void timeout(TimeValue timeout) {
         throw new UnsupportedOperationException("Not supported");
     }