Revert "Delay the creation of SubSearchContext to the FetchSubPhase (#46598)"

This reverts commit 328fe4472fadeb00cbc01aea256167deb7ee88bd.

modules/parent-join/src/main/java/org/elasticsearch/join/query/ParentChildInnerHitContextBuilder.java

@@ -67,26 +67,20 @@ class ParentChildInnerHitContextBuilder extends InnerHitContextBuilder {
     }
 
     @Override
-    public void doValidate(QueryShardContext queryShardContext) {
-        if (ParentJoinFieldMapper.getMapper(queryShardContext.getMapperService()) == null
-                && innerHitBuilder.isIgnoreUnmapped() == false) {
-            throw new IllegalStateException("no join field has been configured");
-        }
-    }
-
-    @Override
-    public void build(SearchContext context, InnerHitsContext innerHitsContext) throws IOException {
+    protected void doBuild(SearchContext context, InnerHitsContext innerHitsContext) throws IOException {
         QueryShardContext queryShardContext = context.getQueryShardContext();
         ParentJoinFieldMapper joinFieldMapper = ParentJoinFieldMapper.getMapper(context.mapperService());
-        if (joinFieldMapper == null) {
-            assert innerHitBuilder.isIgnoreUnmapped() : "should be validated first";
-            return;
+        if (joinFieldMapper != null) {
+            String name = innerHitBuilder.getName() != null ? innerHitBuilder.getName() : typeName;
+            JoinFieldInnerHitSubContext joinFieldInnerHits = new JoinFieldInnerHitSubContext(name, context, typeName,
+                fetchChildInnerHits, joinFieldMapper);
+            setupInnerHitsContext(queryShardContext, joinFieldInnerHits);
+            innerHitsContext.addInnerHitDefinition(joinFieldInnerHits);
+        } else {
+            if (innerHitBuilder.isIgnoreUnmapped() == false) {
+                throw new IllegalStateException("no join field has been configured");
+            }
         }
-        String name = innerHitBuilder.getName() != null ? innerHitBuilder.getName() : typeName;
-        JoinFieldInnerHitSubContext joinFieldInnerHits =
-            new JoinFieldInnerHitSubContext(name, context, typeName, fetchChildInnerHits, joinFieldMapper);
-        setupInnerHitsContext(queryShardContext, joinFieldInnerHits);
-        innerHitsContext.addInnerHitDefinition(joinFieldInnerHits);
     }
 
     static final class JoinFieldInnerHitSubContext extends InnerHitsContext.InnerHitSubContext {
@@ -94,11 +88,8 @@ class ParentChildInnerHitContextBuilder extends InnerHitContextBuilder {
         private final boolean fetchChildInnerHits;
         private final ParentJoinFieldMapper joinFieldMapper;
 
-        JoinFieldInnerHitSubContext(String name,
-                                        SearchContext context,
-                                        String typeName,
-                                        boolean fetchChildInnerHits,
-                                        ParentJoinFieldMapper joinFieldMapper) {
+        JoinFieldInnerHitSubContext(String name, SearchContext context, String typeName, boolean fetchChildInnerHits,
+                                    ParentJoinFieldMapper joinFieldMapper) {
             super(name, context);
             this.typeName = typeName;
             this.fetchChildInnerHits = fetchChildInnerHits;
@@ -111,13 +102,13 @@ class ParentChildInnerHitContextBuilder extends InnerHitContextBuilder {
             TopDocsAndMaxScore[] result = new TopDocsAndMaxScore[hits.length];
             for (int i = 0; i < hits.length; i++) {
                 SearchHit hit = hits[i];
-                String joinName = getSortedDocValue(joinFieldMapper.name(), this, hit.docId());
+                String joinName = getSortedDocValue(joinFieldMapper.name(), context, hit.docId());
                 if (joinName == null) {
                     result[i] = new TopDocsAndMaxScore(Lucene.EMPTY_TOP_DOCS, Float.NaN);
                     continue;
                 }
 
-                QueryShardContext qsc = getQueryShardContext();
+                QueryShardContext qsc = context.getQueryShardContext();
                 ParentIdFieldMapper parentIdFieldMapper =
                     joinFieldMapper.getParentIdFieldMapper(typeName, fetchChildInnerHits == false);
                 if (parentIdFieldMapper == null) {
@@ -135,14 +126,14 @@ class ParentChildInnerHitContextBuilder extends InnerHitContextBuilder {
                         .add(joinFieldMapper.fieldType().termQuery(typeName, qsc), BooleanClause.Occur.FILTER)
                         .build();
                 } else {
-                    String parentId = getSortedDocValue(parentIdFieldMapper.name(), this, hit.docId());
-                    q = mapperService().fullName(IdFieldMapper.NAME).termQuery(parentId, qsc);
+                    String parentId = getSortedDocValue(parentIdFieldMapper.name(), context, hit.docId());
+                    q = context.mapperService().fullName(IdFieldMapper.NAME).termQuery(parentId, qsc);
                 }
 
-                Weight weight = searcher().createWeight(searcher().rewrite(q), ScoreMode.COMPLETE_NO_SCORES, 1f);
+                Weight weight = context.searcher().createWeight(context.searcher().rewrite(q), ScoreMode.COMPLETE_NO_SCORES, 1f);
                 if (size() == 0) {
                     TotalHitCountCollector totalHitCountCollector = new TotalHitCountCollector();
-                    for (LeafReaderContext ctx : searcher().getIndexReader().leaves()) {
+                    for (LeafReaderContext ctx : context.searcher().getIndexReader().leaves()) {
                         intersect(weight, innerHitQueryWeight, totalHitCountCollector, ctx);
                     }
                     result[i] = new TopDocsAndMaxScore(
@@ -151,7 +142,7 @@ class ParentChildInnerHitContextBuilder extends InnerHitContextBuilder {
                             Lucene.EMPTY_SCORE_DOCS
                         ), Float.NaN);
                 } else {
-                    int topN = Math.min(from() + size(), searcher().getIndexReader().maxDoc());
+                    int topN = Math.min(from() + size(), context.searcher().getIndexReader().maxDoc());
                     TopDocsCollector<?> topDocsCollector;
                     MaxScoreCollector maxScoreCollector = null;
                     if (sort() != null) {
@@ -164,7 +155,7 @@ class ParentChildInnerHitContextBuilder extends InnerHitContextBuilder {
                         maxScoreCollector = new MaxScoreCollector();
                     }
                     try {
-                        for (LeafReaderContext ctx : searcher().getIndexReader().leaves()) {
+                        for (LeafReaderContext ctx : context.searcher().getIndexReader().leaves()) {
                             intersect(weight, innerHitQueryWeight, MultiCollector.wrap(topDocsCollector, maxScoreCollector), ctx);
                         }
                     } finally {

modules/parent-join/src/test/java/org/elasticsearch/join/query/HasParentQueryBuilderTests.java

@@ -146,6 +146,7 @@ public class HasParentQueryBuilderTests extends AbstractQueryTestCase<HasParentQ
             // doCreateTestQueryBuilder)
             queryBuilder = (HasParentQueryBuilder) queryBuilder.rewrite(context);
 
+            assertNotNull(context);
             Map<String, InnerHitContextBuilder> innerHitBuilders = new HashMap<>();
             InnerHitContextBuilder.extractInnerHits(queryBuilder, innerHitBuilders);
             assertTrue(innerHitBuilders.containsKey(queryBuilder.innerHit().getName()));

server/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java

@@ -29,6 +29,7 @@ import org.elasticsearch.search.sort.SortAndFormats;
 import org.elasticsearch.search.sort.SortBuilder;
 
 import java.io.IOException;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
 
@@ -46,9 +47,9 @@ public abstract class InnerHitContextBuilder {
         this.query = query;
     }
 
-    public final void validate(QueryShardContext queryShardContext) {
+    public final void build(SearchContext parentSearchContext, InnerHitsContext innerHitsContext) throws IOException {
         long innerResultWindow = innerHitBuilder.getFrom() + innerHitBuilder.getSize();
-        int maxInnerResultWindow = queryShardContext.getIndexSettings().getMaxInnerResultWindow();
+        int maxInnerResultWindow = parentSearchContext.mapperService().getIndexSettings().getMaxInnerResultWindow();
         if (innerResultWindow > maxInnerResultWindow) {
             throw new IllegalArgumentException(
                 "Inner result window is too large, the inner hit definition's [" + innerHitBuilder.getName() +
@@ -57,16 +58,14 @@ public abstract class InnerHitContextBuilder {
                     "] index level setting."
             );
         }
-        doValidate(queryShardContext);
+        doBuild(parentSearchContext, innerHitsContext);
     }
 
     public InnerHitBuilder innerHitBuilder() {
         return innerHitBuilder;
     }
 
-    protected abstract void doValidate(QueryShardContext queryShardContext);
-
-    public abstract void build(SearchContext parentSearchContext, InnerHitsContext innerHitsContext) throws IOException;
+    protected abstract void doBuild(SearchContext parentSearchContext, InnerHitsContext innerHitsContext) throws IOException;
 
     public static void extractInnerHits(QueryBuilder query, Map<String, InnerHitContextBuilder> innerHitBuilders) {
         if (query instanceof AbstractQueryBuilder) {
@@ -114,6 +113,23 @@ public abstract class InnerHitContextBuilder {
         }
         ParsedQuery parsedQuery = new ParsedQuery(query.toQuery(queryShardContext), queryShardContext.copyNamedQueries());
         innerHitsContext.parsedQuery(parsedQuery);
-        innerHitsContext.innerHits(children);
+        Map<String, InnerHitsContext.InnerHitSubContext> baseChildren =
+            buildChildInnerHits(innerHitsContext.parentSearchContext(), children);
+        innerHitsContext.setChildInnerHits(baseChildren);
+    }
+
+    private static Map<String, InnerHitsContext.InnerHitSubContext> buildChildInnerHits(SearchContext parentSearchContext,
+                                Map<String, InnerHitContextBuilder> children) throws IOException {
+
+        Map<String, InnerHitsContext.InnerHitSubContext> childrenInnerHits = new HashMap<>();
+        for (Map.Entry<String, InnerHitContextBuilder> entry : children.entrySet()) {
+            InnerHitsContext childInnerHitsContext = new InnerHitsContext();
+            entry.getValue().build(
+                parentSearchContext, childInnerHitsContext);
+            if (childInnerHitsContext.getInnerHits() != null) {
+                childrenInnerHits.putAll(childInnerHitsContext.getInnerHits());
+            }
+        }
+        return childrenInnerHits;
     }
 }

server/src/main/java/org/elasticsearch/index/query/NestedQueryBuilder.java

@@ -332,34 +332,28 @@ public class NestedQueryBuilder extends AbstractQueryBuilder<NestedQueryBuilder>
     static class NestedInnerHitContextBuilder extends InnerHitContextBuilder {
         private final String path;
 
-        NestedInnerHitContextBuilder(String path,
-                                        QueryBuilder query,
-                                        InnerHitBuilder innerHitBuilder,
-                                        Map<String, InnerHitContextBuilder> children) {
+        NestedInnerHitContextBuilder(String path, QueryBuilder query, InnerHitBuilder innerHitBuilder,
+                                     Map<String, InnerHitContextBuilder> children) {
             super(query, innerHitBuilder, children);
             this.path = path;
         }
 
         @Override
-        public void doValidate(QueryShardContext queryShardContext) {
-            if (queryShardContext.getObjectMapper(path) == null
-                    && innerHitBuilder.isIgnoreUnmapped() == false) {
-                throw new IllegalStateException("[" + query.getName() + "] no mapping found for type [" + path + "]");
-            }
-        }
-
-        @Override
-        public void build(SearchContext searchContext, InnerHitsContext innerHitsContext) throws IOException {
-            QueryShardContext queryShardContext = searchContext.getQueryShardContext();
+        protected void doBuild(SearchContext parentSearchContext,
+                          InnerHitsContext innerHitsContext) throws IOException {
+            QueryShardContext queryShardContext = parentSearchContext.getQueryShardContext();
             ObjectMapper nestedObjectMapper = queryShardContext.getObjectMapper(path);
             if (nestedObjectMapper == null) {
-                assert innerHitBuilder.isIgnoreUnmapped() : "should be validated first";
-                return;
+                if (innerHitBuilder.isIgnoreUnmapped() == false) {
+                    throw new IllegalStateException("[" + query.getName() + "] no mapping found for type [" + path + "]");
+                } else {
+                    return;
+                }
             }
             String name =  innerHitBuilder.getName() != null ? innerHitBuilder.getName() : nestedObjectMapper.fullPath();
             ObjectMapper parentObjectMapper = queryShardContext.nestedScope().nextLevel(nestedObjectMapper);
             NestedInnerHitSubContext nestedInnerHits = new NestedInnerHitSubContext(
-                name, searchContext, parentObjectMapper, nestedObjectMapper
+                name, parentSearchContext, parentObjectMapper, nestedObjectMapper
             );
             setupInnerHitsContext(queryShardContext, nestedInnerHits);
             queryShardContext.nestedScope().previousLevel();
@@ -405,9 +399,9 @@ public class NestedQueryBuilder extends AbstractQueryBuilder<NestedQueryBuilder>
                 LeafReaderContext ctx = searcher().getIndexReader().leaves().get(readerIndex);
 
                 Query childFilter = childObjectMapper.nestedTypeFilter();
-                BitSetProducer parentFilter = bitsetFilterCache().getBitSetProducer(rawParentFilter);
+                BitSetProducer parentFilter = context.bitsetFilterCache().getBitSetProducer(rawParentFilter);
                 Query q = new ParentChildrenBlockJoinQuery(parentFilter, childFilter, parentDocId);
-                Weight weight = searcher().createWeight(searcher().rewrite(q),
+                Weight weight = context.searcher().createWeight(context.searcher().rewrite(q),
                         org.apache.lucene.search.ScoreMode.COMPLETE_NO_SCORES, 1f);
                 if (size() == 0) {
                     TotalHitCountCollector totalHitCountCollector = new TotalHitCountCollector();
@@ -415,7 +409,7 @@ public class NestedQueryBuilder extends AbstractQueryBuilder<NestedQueryBuilder>
                     result[i] = new TopDocsAndMaxScore(new TopDocs(new TotalHits(totalHitCountCollector.getTotalHits(),
                         TotalHits.Relation.EQUAL_TO), Lucene.EMPTY_SCORE_DOCS), Float.NaN);
                 } else {
-                    int topN = Math.min(from() + size(), searcher().getIndexReader().maxDoc());
+                    int topN = Math.min(from() + size(), context.searcher().getIndexReader().maxDoc());
                     TopDocsCollector<?> topDocsCollector;
                     MaxScoreCollector maxScoreCollector = null;
                     if (sort() != null) {

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

@@ -43,7 +43,6 @@ import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.ObjectMapper;
 import org.elasticsearch.index.query.AbstractQueryBuilder;
-import org.elasticsearch.index.query.InnerHitContextBuilder;
 import org.elasticsearch.index.query.ParsedQuery;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryShardContext;
@@ -110,7 +109,6 @@ final class DefaultSearchContext extends SearchContext {
     private ScriptFieldsContext scriptFields;
     private FetchSourceContext fetchSourceContext;
     private DocValueFieldsContext docValueFieldsContext;
-    private Map<String, InnerHitContextBuilder> innerHits = Collections.emptyMap();
     private int from = -1;
     private int size = -1;
     private SortAndFormats sort;
@@ -375,16 +373,6 @@ final class DefaultSearchContext extends SearchContext {
         this.highlight = highlight;
     }
 
-    @Override
-    public void innerHits(Map<String, InnerHitContextBuilder> innerHits) {
-        this.innerHits = innerHits;
-    }
-
-    @Override
-    public Map<String, InnerHitContextBuilder> innerHits() {
-        return innerHits;
-    }
-
     @Override
     public SuggestionSearchContext suggest() {
         return suggest;

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

@@ -739,7 +739,6 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
         context.from(source.from());
         context.size(source.size());
         Map<String, InnerHitContextBuilder> innerHitBuilders = new HashMap<>();
-        context.innerHits(innerHitBuilders);
         if (source.query() != null) {
             InnerHitContextBuilder.extractInnerHits(source.query(), innerHitBuilders);
             context.parsedQuery(queryShardContext.toQuery(source.query()));
@@ -750,7 +749,11 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
         }
         if (innerHitBuilders.size() > 0) {
             for (Map.Entry<String, InnerHitContextBuilder> entry : innerHitBuilders.entrySet()) {
-                entry.getValue().validate(queryShardContext);
+                try {
+                    entry.getValue().build(context, context.innerHits());
+                } catch (IOException e) {
+                    throw new SearchException(shardTarget, "failed to build inner_hits", e);
+                }
             }
         }
         if (source.sorts() != null) {

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

@@ -41,6 +41,7 @@ import java.io.IOException;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Objects;
 
 /**
  * Context used for inner hits retrieval
@@ -52,6 +53,10 @@ public final class InnerHitsContext {
         this.innerHits = new HashMap<>();
     }
 
+    InnerHitsContext(Map<String, InnerHitSubContext> innerHits) {
+        this.innerHits = Objects.requireNonNull(innerHits);
+    }
+
     public Map<String, InnerHitSubContext> getInnerHits() {
         return innerHits;
     }
@@ -72,6 +77,8 @@ public final class InnerHitsContext {
     public abstract static class InnerHitSubContext extends SubSearchContext {
 
         private final String name;
+        protected final SearchContext context;
+        private InnerHitsContext childInnerHits;
 
         // TODO: when types are complete removed just use String instead for the id:
         private Uid uid;
@@ -79,6 +86,7 @@ public final class InnerHitsContext {
         protected InnerHitSubContext(String name, SearchContext context) {
             super(context);
             this.name = name;
+            this.context = context;
         }
 
         public abstract TopDocsAndMaxScore[] topDocs(SearchHit[] hits) throws IOException;
@@ -87,12 +95,25 @@ public final class InnerHitsContext {
             return name;
         }
 
+        @Override
+        public InnerHitsContext innerHits() {
+            return childInnerHits;
+        }
+
+        public void setChildInnerHits(Map<String, InnerHitSubContext> childInnerHits) {
+            this.childInnerHits = new InnerHitsContext(childInnerHits);
+        }
+
         protected Weight createInnerHitQueryWeight() throws IOException {
             final boolean needsScores = size() != 0 && (sort() == null || sort().sort.needsScores());
-            return searcher().createWeight(searcher().rewrite(query()),
+            return context.searcher().createWeight(context.searcher().rewrite(query()),
                     needsScores ? ScoreMode.COMPLETE : ScoreMode.COMPLETE_NO_SCORES, 1f);
         }
 
+        public SearchContext parentSearchContext() {
+            return context;
+        }
+
         public Uid getUid() {
             return uid;
         }

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

@@ -24,7 +24,6 @@ import org.apache.lucene.search.ScoreDoc;
 import org.elasticsearch.common.lucene.search.TopDocsAndMaxScore;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.Uid;
-import org.elasticsearch.index.query.InnerHitContextBuilder;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.SearchHits;
 import org.elasticsearch.search.fetch.FetchPhase;
@@ -46,16 +45,11 @@ public final class InnerHitsFetchSubPhase implements FetchSubPhase {
 
     @Override
     public void hitsExecute(SearchContext context, SearchHit[] hits) throws IOException {
-        if (context.innerHits().isEmpty()) {
+        if ((context.innerHits() != null && context.innerHits().getInnerHits().size() > 0) == false) {
             return;
         }
 
-        final InnerHitsContext innerHitsContext = new InnerHitsContext();
-        for (Map.Entry<String, InnerHitContextBuilder> entry : context.innerHits().entrySet()) {
-            entry.getValue().build(context, innerHitsContext);
-        }
-
-        for (Map.Entry<String, InnerHitsContext.InnerHitSubContext> entry : innerHitsContext.getInnerHits().entrySet()) {
+        for (Map.Entry<String, InnerHitsContext.InnerHitSubContext> entry : context.innerHits().getInnerHits().entrySet()) {
             InnerHitsContext.InnerHitSubContext innerHits = entry.getValue();
             TopDocsAndMaxScore[] topDocs = innerHits.topDocs(hits);
             for (int i = 0; i < hits.length; i++) {

server/src/main/java/org/elasticsearch/search/internal/FilteredSearchContext.java

@@ -31,7 +31,6 @@ import org.elasticsearch.index.fielddata.IndexFieldData;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.ObjectMapper;
-import org.elasticsearch.index.query.InnerHitContextBuilder;
 import org.elasticsearch.index.query.ParsedQuery;
 import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.index.shard.IndexShard;
@@ -45,6 +44,7 @@ import org.elasticsearch.search.fetch.FetchPhase;
 import org.elasticsearch.search.fetch.FetchSearchResult;
 import org.elasticsearch.search.fetch.StoredFieldsContext;
 import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
+import org.elasticsearch.search.fetch.subphase.InnerHitsContext;
 import org.elasticsearch.search.fetch.subphase.ScriptFieldsContext;
 import org.elasticsearch.search.fetch.subphase.highlight.SearchContextHighlight;
 import org.elasticsearch.search.lookup.SearchLookup;
@@ -176,15 +176,10 @@ public abstract class FilteredSearchContext extends SearchContext {
     }
 
     @Override
-    public Map<String, InnerHitContextBuilder> innerHits() {
+    public InnerHitsContext innerHits() {
         return in.innerHits();
     }
 
-    @Override
-    public void innerHits(Map<String, InnerHitContextBuilder> innerHits) {
-        in.innerHits(innerHits);
-    }
-
     @Override
     public SuggestionSearchContext suggest() {
         return in.suggest();

server/src/main/java/org/elasticsearch/search/internal/SearchContext.java

@@ -18,6 +18,7 @@
  */
 package org.elasticsearch.search.internal;
 
+
 import org.apache.lucene.search.Collector;
 import org.apache.lucene.search.FieldDoc;
 import org.apache.lucene.search.Query;
@@ -36,7 +37,6 @@ import org.elasticsearch.index.fielddata.IndexFieldData;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.ObjectMapper;
-import org.elasticsearch.index.query.InnerHitContextBuilder;
 import org.elasticsearch.index.query.ParsedQuery;
 import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.index.shard.IndexShard;
@@ -51,6 +51,7 @@ import org.elasticsearch.search.fetch.FetchSearchResult;
 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;
 import org.elasticsearch.search.fetch.subphase.highlight.SearchContextHighlight;
 import org.elasticsearch.search.lookup.SearchLookup;
@@ -86,6 +87,7 @@ public abstract class SearchContext extends AbstractRefCounted implements Releas
 
     private Map<Lifetime, List<Releasable>> clearables = null;
     private final AtomicBoolean closed = new AtomicBoolean(false);
+    private InnerHitsContext innerHitsContext;
 
     protected SearchContext() {
         super("search_context");
@@ -162,9 +164,12 @@ public abstract class SearchContext extends AbstractRefCounted implements Releas
 
     public abstract void highlight(SearchContextHighlight highlight);
 
-    public abstract void innerHits(Map<String, InnerHitContextBuilder> innerHits);
-
-    public abstract Map<String, InnerHitContextBuilder> innerHits();
+    public InnerHitsContext innerHits() {
+        if (innerHitsContext == null) {
+            innerHitsContext = new InnerHitsContext();
+        }
+        return innerHitsContext;
+    }
 
     public abstract SuggestionSearchContext suggest();
 

server/src/main/java/org/elasticsearch/search/internal/SubSearchContext.java

@@ -20,9 +20,7 @@ package org.elasticsearch.search.internal;
 
 import org.apache.lucene.search.Query;
 import org.elasticsearch.common.unit.TimeValue;
-import org.elasticsearch.index.query.InnerHitContextBuilder;
 import org.elasticsearch.index.query.ParsedQuery;
-import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.search.aggregations.SearchContextAggregations;
 import org.elasticsearch.search.collapse.CollapseContext;
 import org.elasticsearch.search.fetch.FetchSearchResult;
@@ -31,15 +29,12 @@ 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;
-import org.elasticsearch.search.lookup.SearchLookup;
 import org.elasticsearch.search.query.QuerySearchResult;
 import org.elasticsearch.search.rescore.RescoreContext;
 import org.elasticsearch.search.sort.SortAndFormats;
 import org.elasticsearch.search.suggest.SuggestionSearchContext;
 
-import java.util.Collections;
 import java.util.List;
-import java.util.Map;
 
 public class SubSearchContext extends FilteredSearchContext {
 
@@ -47,8 +42,6 @@ public class SubSearchContext extends FilteredSearchContext {
     // the to hits are returned per bucket.
     private static final int DEFAULT_SIZE = 3;
 
-    private final QueryShardContext queryShardContext;
-
     private int from;
     private int size = DEFAULT_SIZE;
     private SortAndFormats sort;
@@ -67,7 +60,6 @@ public class SubSearchContext extends FilteredSearchContext {
     private FetchSourceContext fetchSourceContext;
     private DocValueFieldsContext docValueFieldsContext;
     private SearchContextHighlight highlight;
-    private Map<String, InnerHitContextBuilder> innerHits = Collections.emptyMap();
 
     private boolean explain;
     private boolean trackScores;
@@ -78,9 +70,6 @@ public class SubSearchContext extends FilteredSearchContext {
         super(context);
         this.fetchSearchResult = new FetchSearchResult();
         this.querySearchResult = new QuerySearchResult();
-        // we clone the query shard context in the sub context because the original one
-        // might be frozen at this point.
-        this.queryShardContext = new QueryShardContext(context.getQueryShardContext());
     }
 
     @Override
@@ -91,11 +80,6 @@ public class SubSearchContext extends FilteredSearchContext {
     public void preProcess(boolean rewrite) {
     }
 
-    @Override
-    public QueryShardContext getQueryShardContext() {
-        return queryShardContext;
-    }
-
     @Override
     public Query buildFilteredQuery(Query query) {
         throw new UnsupportedOperationException("this context should be read only");
@@ -372,19 +356,4 @@ public class SubSearchContext extends FilteredSearchContext {
     public long getRelativeTimeInMillis() {
         throw new UnsupportedOperationException("Not supported");
     }
-
-    @Override
-    public Map<String, InnerHitContextBuilder> innerHits() {
-        return innerHits;
-    }
-
-    @Override
-    public void innerHits(Map<String, InnerHitContextBuilder> innerHits) {
-        this.innerHits = innerHits;
-    }
-
-    @Override
-    public SearchLookup lookup() {
-        return queryShardContext.lookup();
-    }
 }

server/src/test/java/org/elasticsearch/index/query/NestedQueryBuilderTests.java

@@ -321,7 +321,7 @@ public class NestedQueryBuilderTests extends AbstractQueryTestCase<NestedQueryBu
 
         MapperService mapperService = mock(MapperService.class);
         IndexSettings settings = new IndexSettings(newIndexMeta("index", Settings.EMPTY), Settings.EMPTY);
-        when(queryShardContext.getIndexSettings()).thenReturn(settings);
+        when(mapperService.getIndexSettings()).thenReturn(settings);
         when(searchContext.mapperService()).thenReturn(mapperService);
 
         InnerHitBuilder leafInnerHits = randomNestedInnerHits();
@@ -333,7 +333,7 @@ public class NestedQueryBuilderTests extends AbstractQueryTestCase<NestedQueryBu
             query1.extractInnerHitBuilders(innerHitBuilders);
             assertThat(innerHitBuilders.size(), Matchers.equalTo(1));
             assertTrue(innerHitBuilders.containsKey(leafInnerHits.getName()));
-            innerHitBuilders.get(leafInnerHits.getName()).validate(searchContext.getQueryShardContext());
+            innerHitBuilders.get(leafInnerHits.getName()).build(searchContext, innerHitsContext);
         });
         innerHitBuilders.clear();
         NestedQueryBuilder query2 = new NestedQueryBuilder("path", new MatchAllQueryBuilder(), ScoreMode.None);

test/framework/src/main/java/org/elasticsearch/test/TestSearchContext.java

@@ -32,7 +32,6 @@ import org.elasticsearch.index.fielddata.IndexFieldData;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.ObjectMapper;
-import org.elasticsearch.index.query.InnerHitContextBuilder;
 import org.elasticsearch.index.query.ParsedQuery;
 import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.index.shard.IndexShard;
@@ -204,14 +203,6 @@ public class TestSearchContext extends SearchContext {
     public void highlight(SearchContextHighlight highlight) {
     }
 
-    @Override
-    public void innerHits(Map<String, InnerHitContextBuilder> innerHits) {}
-
-    @Override
-    public Map<String, InnerHitContextBuilder> innerHits() {
-        return null;
-    }
-
     @Override
     public SuggestionSearchContext suggest() {
         return null;