refactor sub fetch phase to also allow for hits level execution

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

@@ -28,10 +28,10 @@ import org.elasticsearch.search.controller.SearchPhaseController;
 import org.elasticsearch.search.dfs.DfsPhase;
 import org.elasticsearch.search.facet.FacetModule;
 import org.elasticsearch.search.fetch.FetchPhase;
-import org.elasticsearch.search.fetch.explain.ExplainSearchHitPhase;
+import org.elasticsearch.search.fetch.explain.ExplainFetchSubPhase;
-import org.elasticsearch.search.fetch.matchedfilters.MatchedFiltersSearchHitPhase;
+import org.elasticsearch.search.fetch.matchedfilters.MatchedFiltersFetchSubPhase;
-import org.elasticsearch.search.fetch.script.ScriptFieldsSearchHitPhase;
+import org.elasticsearch.search.fetch.script.ScriptFieldsFetchSubPhase;
-import org.elasticsearch.search.fetch.version.VersionSearchHitPhase;
+import org.elasticsearch.search.fetch.version.VersionFetchSubPhase;
 import org.elasticsearch.search.highlight.HighlightPhase;
 import org.elasticsearch.search.query.QueryPhase;
 
@@ -51,10 +51,10 @@ public class SearchModule extends AbstractModule implements SpawnModules {
         bind(SearchPhaseController.class).asEagerSingleton();
 
         bind(FetchPhase.class).asEagerSingleton();
-        bind(ExplainSearchHitPhase.class).asEagerSingleton();
+        bind(ExplainFetchSubPhase.class).asEagerSingleton();
-        bind(ScriptFieldsSearchHitPhase.class).asEagerSingleton();
+        bind(ScriptFieldsFetchSubPhase.class).asEagerSingleton();
-        bind(VersionSearchHitPhase.class).asEagerSingleton();
+        bind(VersionFetchSubPhase.class).asEagerSingleton();
-        bind(MatchedFiltersSearchHitPhase.class).asEagerSingleton();
+        bind(MatchedFiltersFetchSubPhase.class).asEagerSingleton();
         bind(HighlightPhase.class).asEagerSingleton();
 
         bind(SearchServiceTransportAction.class).asEagerSingleton();

modules/elasticsearch/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java

@@ -41,10 +41,10 @@ import org.elasticsearch.indices.TypeMissingException;
 import org.elasticsearch.search.SearchHitField;
 import org.elasticsearch.search.SearchParseElement;
 import org.elasticsearch.search.SearchPhase;
-import org.elasticsearch.search.fetch.explain.ExplainSearchHitPhase;
+import org.elasticsearch.search.fetch.explain.ExplainFetchSubPhase;
-import org.elasticsearch.search.fetch.matchedfilters.MatchedFiltersSearchHitPhase;
+import org.elasticsearch.search.fetch.matchedfilters.MatchedFiltersFetchSubPhase;
-import org.elasticsearch.search.fetch.script.ScriptFieldsSearchHitPhase;
+import org.elasticsearch.search.fetch.script.ScriptFieldsFetchSubPhase;
-import org.elasticsearch.search.fetch.version.VersionSearchHitPhase;
+import org.elasticsearch.search.fetch.version.VersionFetchSubPhase;
 import org.elasticsearch.search.highlight.HighlightPhase;
 import org.elasticsearch.search.internal.InternalSearchHit;
 import org.elasticsearch.search.internal.InternalSearchHitField;
@@ -62,18 +62,18 @@ import java.util.Map;
  */
 public class FetchPhase implements SearchPhase {
 
-    private final SearchHitPhase[] hitPhases;
+    private final FetchSubPhase[] fetchSubPhases;
 
-    @Inject public FetchPhase(HighlightPhase highlightPhase, ScriptFieldsSearchHitPhase scriptFieldsPhase,
+    @Inject public FetchPhase(HighlightPhase highlightPhase, ScriptFieldsFetchSubPhase scriptFieldsPhase,
-                              MatchedFiltersSearchHitPhase matchFiltersPhase, ExplainSearchHitPhase explainPhase, VersionSearchHitPhase versionPhase) {
+                              MatchedFiltersFetchSubPhase matchFiltersPhase, ExplainFetchSubPhase explainPhase, VersionFetchSubPhase versionPhase) {
-        this.hitPhases = new SearchHitPhase[]{scriptFieldsPhase, matchFiltersPhase, explainPhase, highlightPhase, versionPhase};
+        this.fetchSubPhases = new FetchSubPhase[]{scriptFieldsPhase, matchFiltersPhase, explainPhase, highlightPhase, versionPhase};
     }
 
     @Override public Map<String, ? extends SearchParseElement> parseElements() {
         ImmutableMap.Builder<String, SearchParseElement> parseElements = ImmutableMap.builder();
         parseElements.put("fields", new FieldsParseElement());
-        for (SearchHitPhase hitPhase : hitPhases) {
+        for (FetchSubPhase fetchSubPhase : fetchSubPhases) {
-            parseElements.putAll(hitPhase.parseElements());
+            parseElements.putAll(fetchSubPhase.parseElements());
         }
         return parseElements.build();
     }
@@ -199,14 +199,21 @@ public class FetchPhase implements SearchPhase {
                 }
             }
 
-            for (SearchHitPhase hitPhase : hitPhases) {
+            for (FetchSubPhase fetchSubPhase : fetchSubPhases) {
-                SearchHitPhase.HitContext hitContext = new SearchHitPhase.HitContext();
+                FetchSubPhase.HitContext hitContext = new FetchSubPhase.HitContext();
-                if (hitPhase.executionNeeded(context)) {
+                if (fetchSubPhase.hitExecutionNeeded(context)) {
                     hitContext.reset(searchHit, subReader, subDoc, doc);
-                    hitPhase.execute(context, hitContext);
+                    fetchSubPhase.hitExecute(context, hitContext);
                 }
             }
         }
+
+        for (FetchSubPhase fetchSubPhase : fetchSubPhases) {
+            if (fetchSubPhase.hitsExecutionNeeded(context)) {
+                fetchSubPhase.hitsExecute(context, hits);
+            }
+        }
+
         context.fetchResult().hits(new InternalSearchHits(hits, context.queryResult().topDocs().totalHits, context.queryResult().topDocs().getMaxScore()));
     }
 

modules/elasticsearch/src/main/java/org/elasticsearch/search/fetch/FetchSubPhase.java

@@ -22,7 +22,6 @@ package org.elasticsearch.search.fetch;
 import org.apache.lucene.document.Document;
 import org.apache.lucene.index.IndexReader;
 import org.elasticsearch.ElasticSearchException;
-import org.elasticsearch.index.mapper.Uid;
 import org.elasticsearch.search.SearchParseElement;
 import org.elasticsearch.search.internal.InternalSearchHit;
 import org.elasticsearch.search.internal.SearchContext;
@@ -32,7 +31,7 @@ import java.util.Map;
 /**
  * @author kimchy (shay.banon)
  */
-public interface SearchHitPhase {
+public interface FetchSubPhase {
 
     public static class HitContext {
         private InternalSearchHit hit;
@@ -66,10 +65,14 @@ public interface SearchHitPhase {
 
     Map<String, ? extends SearchParseElement> parseElements();
 
-    boolean executionNeeded(SearchContext context);
+    boolean hitExecutionNeeded(SearchContext context);
 
     /**
      * Executes the hit level phase, with a reader and doc id (note, its a low level reader, and the matching doc).
      */
-    void execute(SearchContext context, HitContext hitContext) throws ElasticSearchException;
+    void hitExecute(SearchContext context, HitContext hitContext) throws ElasticSearchException;
+
+    boolean hitsExecutionNeeded(SearchContext context);
+
+    void hitsExecute(SearchContext context, InternalSearchHit[] hits) throws ElasticSearchException;
 }

modules/elasticsearch/src/main/java/org/elasticsearch/search/fetch/explain/ExplainFetchSubPhase.java

@@ -23,7 +23,8 @@ import org.elasticsearch.ElasticSearchException;
 import org.elasticsearch.common.collect.ImmutableMap;
 import org.elasticsearch.search.SearchParseElement;
 import org.elasticsearch.search.fetch.FetchPhaseExecutionException;
-import org.elasticsearch.search.fetch.SearchHitPhase;
+import org.elasticsearch.search.fetch.FetchSubPhase;
+import org.elasticsearch.search.internal.InternalSearchHit;
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
@@ -32,17 +33,24 @@ import java.util.Map;
 /**
  * @author kimchy (shay.banon)
  */
-public class ExplainSearchHitPhase implements SearchHitPhase {
+public class ExplainFetchSubPhase implements FetchSubPhase {
 
     @Override public Map<String, ? extends SearchParseElement> parseElements() {
         return ImmutableMap.of("explain", new ExplainParseElement());
     }
 
-    @Override public boolean executionNeeded(SearchContext context) {
+    @Override public boolean hitsExecutionNeeded(SearchContext context) {
+        return false;
+    }
+
+    @Override public void hitsExecute(SearchContext context, InternalSearchHit[] hits) throws ElasticSearchException {
+    }
+
+    @Override public boolean hitExecutionNeeded(SearchContext context) {
         return context.explain();
     }
 
-    @Override public void execute(SearchContext context, HitContext hitContext) throws ElasticSearchException {
+    @Override public void hitExecute(SearchContext context, HitContext hitContext) throws ElasticSearchException {
         try {
             // we use the top level doc id, since we work with the top level searcher
             hitContext.hit().explanation(context.searcher().explain(context.query(), hitContext.hit().docId()));

modules/elasticsearch/src/main/java/org/elasticsearch/search/fetch/matchedfilters/MatchedFiltersFetchSubPhase.java

@@ -19,7 +19,6 @@
 
 package org.elasticsearch.search.fetch.matchedfilters;
 
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.search.DocIdSet;
 import org.apache.lucene.search.Filter;
 import org.elasticsearch.ElasticSearchException;
@@ -27,9 +26,8 @@ import org.elasticsearch.common.collect.ImmutableMap;
 import org.elasticsearch.common.collect.Lists;
 import org.elasticsearch.common.lucene.docset.DocSet;
 import org.elasticsearch.common.lucene.docset.DocSets;
-import org.elasticsearch.index.mapper.Uid;
 import org.elasticsearch.search.SearchParseElement;
-import org.elasticsearch.search.fetch.SearchHitPhase;
+import org.elasticsearch.search.fetch.FetchSubPhase;
 import org.elasticsearch.search.internal.InternalSearchHit;
 import org.elasticsearch.search.internal.SearchContext;
 
@@ -40,17 +38,24 @@ import java.util.Map;
 /**
  * @author kimchy (shay.banon)
  */
-public class MatchedFiltersSearchHitPhase implements SearchHitPhase {
+public class MatchedFiltersFetchSubPhase implements FetchSubPhase {
 
     @Override public Map<String, ? extends SearchParseElement> parseElements() {
         return ImmutableMap.of();
     }
 
-    @Override public boolean executionNeeded(SearchContext context) {
+    @Override public boolean hitsExecutionNeeded(SearchContext context) {
+        return false;
+    }
+
+    @Override public void hitsExecute(SearchContext context, InternalSearchHit[] hits) throws ElasticSearchException {
+    }
+
+    @Override public boolean hitExecutionNeeded(SearchContext context) {
         return !context.parsedQuery().namedFilters().isEmpty();
     }
 
-    @Override public void execute(SearchContext context, HitContext hitContext) throws ElasticSearchException {
+    @Override public void hitExecute(SearchContext context, HitContext hitContext) throws ElasticSearchException {
         List<String> matchedFilters = Lists.newArrayListWithCapacity(2);
         for (Map.Entry<String, Filter> entry : context.parsedQuery().namedFilters().entrySet()) {
             String name = entry.getKey();

modules/elasticsearch/src/main/java/org/elasticsearch/search/fetch/script/ScriptFieldsFetchSubPhase.java

@@ -24,7 +24,8 @@ import org.elasticsearch.common.collect.ImmutableMap;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.search.SearchHitField;
 import org.elasticsearch.search.SearchParseElement;
-import org.elasticsearch.search.fetch.SearchHitPhase;
+import org.elasticsearch.search.fetch.FetchSubPhase;
+import org.elasticsearch.search.internal.InternalSearchHit;
 import org.elasticsearch.search.internal.InternalSearchHitField;
 import org.elasticsearch.search.internal.SearchContext;
 
@@ -35,9 +36,9 @@ import java.util.Map;
 /**
  * @author kimchy (shay.banon)
  */
-public class ScriptFieldsSearchHitPhase implements SearchHitPhase {
+public class ScriptFieldsFetchSubPhase implements FetchSubPhase {
 
-    @Inject public ScriptFieldsSearchHitPhase() {
+    @Inject public ScriptFieldsFetchSubPhase() {
     }
 
     @Override public Map<String, ? extends SearchParseElement> parseElements() {
@@ -47,11 +48,18 @@ public class ScriptFieldsSearchHitPhase implements SearchHitPhase {
         return parseElements.build();
     }
 
-    @Override public boolean executionNeeded(SearchContext context) {
+    @Override public boolean hitsExecutionNeeded(SearchContext context) {
+        return false;
+    }
+
+    @Override public void hitsExecute(SearchContext context, InternalSearchHit[] hits) throws ElasticSearchException {
+    }
+
+    @Override public boolean hitExecutionNeeded(SearchContext context) {
         return context.hasScriptFields();
     }
 
-    @Override public void execute(SearchContext context, HitContext hitContext) throws ElasticSearchException {
+    @Override public void hitExecute(SearchContext context, HitContext hitContext) throws ElasticSearchException {
         for (ScriptFieldsContext.ScriptField scriptField : context.scriptFields().fields()) {
             scriptField.script().setNextReader(hitContext.reader());
             scriptField.script().setNextDocId(hitContext.docId());

modules/elasticsearch/src/main/java/org/elasticsearch/search/fetch/version/VersionFetchSubPhase.java

@@ -24,7 +24,8 @@ import org.elasticsearch.common.collect.ImmutableMap;
 import org.elasticsearch.common.lucene.uid.UidField;
 import org.elasticsearch.index.mapper.internal.UidFieldMapper;
 import org.elasticsearch.search.SearchParseElement;
-import org.elasticsearch.search.fetch.SearchHitPhase;
+import org.elasticsearch.search.fetch.FetchSubPhase;
+import org.elasticsearch.search.internal.InternalSearchHit;
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.util.Map;
@@ -32,17 +33,24 @@ import java.util.Map;
 /**
  * @author kimchy (shay.banon)
  */
-public class VersionSearchHitPhase implements SearchHitPhase {
+public class VersionFetchSubPhase implements FetchSubPhase {
 
     @Override public Map<String, ? extends SearchParseElement> parseElements() {
         return ImmutableMap.of("version", new VersionParseElement());
     }
 
-    @Override public boolean executionNeeded(SearchContext context) {
+    @Override public boolean hitsExecutionNeeded(SearchContext context) {
+        return false;
+    }
+
+    @Override public void hitsExecute(SearchContext context, InternalSearchHit[] hits) throws ElasticSearchException {
+    }
+
+    @Override public boolean hitExecutionNeeded(SearchContext context) {
         return context.version();
     }
 
-    @Override public void execute(SearchContext context, HitContext hitContext) throws ElasticSearchException {
+    @Override public void hitExecute(SearchContext context, HitContext hitContext) throws ElasticSearchException {
         // it might make sense to cache the TermDocs on a shared fetch context and just skip here)
         // it is going to mean we work on the high level multi reader and not the lower level reader as is
         // the case below...

modules/elasticsearch/src/main/java/org/elasticsearch/search/highlight/HighlightPhase.java

@@ -40,9 +40,10 @@ import org.elasticsearch.index.mapper.FieldMapper;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.search.SearchParseElement;
 import org.elasticsearch.search.fetch.FetchPhaseExecutionException;
-import org.elasticsearch.search.fetch.SearchHitPhase;
+import org.elasticsearch.search.fetch.FetchSubPhase;
 import org.elasticsearch.search.highlight.vectorhighlight.SourceScoreOrderFragmentsBuilder;
 import org.elasticsearch.search.highlight.vectorhighlight.SourceSimpleFragmentsBuilder;
+import org.elasticsearch.search.internal.InternalSearchHit;
 import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.lookup.SearchLookup;
 
@@ -58,7 +59,7 @@ import static org.elasticsearch.common.collect.Maps.*;
 /**
  * @author kimchy (shay.banon)
  */
-public class HighlightPhase implements SearchHitPhase {
+public class HighlightPhase implements FetchSubPhase {
 
     public static class Encoders {
         public static Encoder DEFAULT = new DefaultEncoder();
@@ -69,11 +70,18 @@ public class HighlightPhase implements SearchHitPhase {
         return ImmutableMap.of("highlight", new HighlighterParseElement());
     }
 
-    @Override public boolean executionNeeded(SearchContext context) {
+    @Override public boolean hitsExecutionNeeded(SearchContext context) {
+        return false;
+    }
+
+    @Override public void hitsExecute(SearchContext context, InternalSearchHit[] hits) throws ElasticSearchException {
+    }
+
+    @Override public boolean hitExecutionNeeded(SearchContext context) {
         return context.highlight() != null;
     }
 
-    @Override public void execute(SearchContext context, HitContext hitContext) throws ElasticSearchException {
+    @Override public void hitExecute(SearchContext context, HitContext hitContext) throws ElasticSearchException {
         try {
             DocumentMapper documentMapper = context.mapperService().documentMapper(hitContext.hit().type());
 

modules/elasticsearch/src/main/java/org/elasticsearch/search/internal/InternalSearchHit.java

@@ -69,7 +69,7 @@ public class InternalSearchHit implements SearchHit {
 
     private Map<String, SearchHitField> fields = ImmutableMap.of();
 
-    private Map<String, HighlightField> highlightFields = ImmutableMap.of();
+    private Map<String, HighlightField> highlightFields = null;
 
     private Object[] sortValues = EMPTY_SORT_VALUES;
 
@@ -230,7 +230,14 @@ public class InternalSearchHit implements SearchHit {
         this.fields = fields;
     }
 
+    public Map<String, HighlightField> internalHighlightFields() {
+        return highlightFields;
+    }
+
     @Override public Map<String, HighlightField> highlightFields() {
+        if (highlightFields == null) {
+            return ImmutableMap.of();
+        }
         return this.highlightFields;
     }