Use new search lookup for each inner hits context

Author: romseygeek

server/src/internalClusterTest/java/org/elasticsearch/search/fetch/subphase/InnerHitsIT.java

@@ -8,7 +8,6 @@
 
 package org.elasticsearch.search.fetch.subphase;
 
-import com.carrotsearch.randomizedtesting.annotations.Seed;
 import org.apache.lucene.search.join.ScoreMode;
 import org.apache.lucene.util.ArrayUtil;
 import org.elasticsearch.action.index.IndexRequestBuilder;
@@ -61,14 +60,18 @@
 import static org.hamcrest.Matchers.notNullValue;
 import static org.hamcrest.Matchers.nullValue;
 
-@Seed("9E43F62788D47F1E:E99EDE15224E9745")
 public class InnerHitsIT extends ESIntegTestCase {
 
     @Override
     protected Collection<Class<? extends Plugin>> nodePlugins() {
         return Arrays.asList(InternalSettingsPlugin.class, CustomScriptPlugin.class);
     }
 
+    @Override
+    protected int numberOfShards() {
+        return 1;
+    }
+
     public static class CustomScriptPlugin extends MockScriptPlugin {
         @Override
         protected Map<String, Function<Map<String, Object>, Object>> pluginScripts() {
@@ -197,12 +200,12 @@ public void testRandomNested() throws Exception {
                     .field("foo", i)
                     .startArray("field1");
             for (int j = 0; j < numInnerObjects; j++) {
-                source.startObject().field("x", "y").endObject();
+                source.startObject().field("x", "y" + i + ":" + j).endObject();
             }
             numInnerObjects = field2InnerObjects[i] = scaledRandomIntBetween(1, numDocs);
             source.endArray().startArray("field2");
             for (int j = 0; j < numInnerObjects; j++) {
-                source.startObject().field("x", "y").endObject();
+                source.startObject().field("p", "q" + i + ":" + j).endObject();
             }
             source.endArray().endObject();
             requestBuilders.add(client().prepareIndex("idx").setId(Integer.toString(i)).setSource(source));

server/src/main/java/org/elasticsearch/search/fetch/FetchContext.java

@@ -22,7 +22,6 @@
 import org.elasticsearch.search.fetch.subphase.highlight.SearchHighlightContext;
 import org.elasticsearch.search.internal.ContextIndexSearcher;
 import org.elasticsearch.search.internal.SearchContext;
-import org.elasticsearch.search.lookup.SearchLookup;
 import org.elasticsearch.search.lookup.SourceLookup;
 import org.elasticsearch.search.rescore.RescoreContext;
 
@@ -35,14 +34,12 @@
 public class FetchContext {
 
     private final SearchContext searchContext;
-    private final SearchLookup searchLookup;
 
     /**
      * Create a FetchContext based on a SearchContext
      */
     public FetchContext(SearchContext searchContext) {
         this.searchContext = searchContext;
-        this.searchLookup = searchContext.getSearchExecutionContext().lookup();
     }
 
     /**
@@ -59,13 +56,6 @@ public ContextIndexSearcher searcher() {
         return searchContext.searcher();
     }
 
-    /**
-     * The {@code SearchLookup} for the this context
-     */
-    public SearchLookup searchLookup() {
-        return searchLookup;
-    }
-
     /**
      * The original query
      */

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

@@ -37,6 +37,7 @@
 import org.elasticsearch.search.fetch.subphase.InnerHitsPhase;
 import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.lookup.LeafSearchLookup;
+import org.elasticsearch.search.lookup.SearchLookup;
 import org.elasticsearch.search.lookup.SourceLookup;
 import org.elasticsearch.search.lookup.ValuesLookup;
 import org.elasticsearch.tasks.TaskCancelledException;
@@ -102,6 +103,7 @@ public void execute(SearchContext context) {
 
         List<FetchSubPhaseProcessor> processors = getProcessors(context.shardTarget(), fetchContext);
         NestedDocuments nestedDocuments = context.getSearchExecutionContext().getNestedDocuments();
+        SearchLookup searchLookup = context.getSearchLookup();
 
         int currentReaderIndex = -1;
         LeafReaderContext currentReaderContext = null;
@@ -134,11 +136,11 @@ public void execute(SearchContext context) {
                         processor.setNextReader(currentReaderContext);
                     }
                     leafNestedDocuments = nestedDocuments.getLeafNestedDocuments(currentReaderContext);
-                    currentValuesLookup = fetchContext.searchLookup().getLeafSearchLookup(currentReaderContext);
+                    currentValuesLookup = searchLookup.getLeafSearchLookup(currentReaderContext);
                 }
                 assert currentReaderContext != null;
                 assert currentValuesLookup != null;
-                currentValuesLookup.setDocument(docId);
+                currentValuesLookup.setDocument(docId - currentReaderContext.docBase);
                 HitContext hit = prepareHitContext(
                     context,
                     leafNestedDocuments,
@@ -377,9 +379,6 @@ private HitContext prepareNestedHitContext(SearchContext context,
 
         SearchHit.NestedIdentity nestedIdentity = nestedInfo.nestedIdentity();
 
-        SearchHit hit = new SearchHit(topDocId, rootId, nestedIdentity, docFields, metaFields);
-        HitContext hitContext = new HitContext(hit, valuesLookup, subReaderContext, nestedInfo.doc());
-
         if (rootSourceAsMap != null && rootSourceAsMap.isEmpty() == false) {
             // Isolate the nested json array object that matches with nested hit and wrap it back into the same json
             // structure with the nested json array object being the actual content. The latter is important, so that
@@ -413,10 +412,10 @@ private HitContext prepareNestedHitContext(SearchContext context,
                 }
             }
 
-            hitContext.valuesLookup().source().setSource(nestedSourceAsMap);
-            hitContext.valuesLookup().source().setSourceContentType(rootSourceContentType);
+            valuesLookup = ValuesLookup.wrapWithSource(valuesLookup, nestedSourceAsMap, rootSourceContentType);
         }
-        return hitContext;
+        SearchHit hit = new SearchHit(topDocId, rootId, nestedIdentity, docFields, metaFields);
+        return new HitContext(hit, valuesLookup, subReaderContext, nestedInfo.doc());
     }
 
     private void loadStoredFields(Function<String, MappedFieldType> fieldTypeLookup,

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

@@ -21,9 +21,11 @@
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.util.Bits;
 import org.elasticsearch.common.lucene.search.TopDocsAndMaxScore;
+import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.internal.SubSearchContext;
+import org.elasticsearch.search.lookup.SearchLookup;
 import org.elasticsearch.search.lookup.SourceLookup;
 
 import java.io.IOException;
@@ -69,6 +71,7 @@ public abstract static class InnerHitSubContext extends SubSearchContext {
         protected final SearchContext context;
         private InnerHitsContext childInnerHits;
         private Weight innerHitQueryWeight;
+        private SearchLookup searchLookup;
 
         private String rootId;
         private final SourceLookup rootLookup = new SourceLookup();
@@ -134,6 +137,15 @@ public void setRootLookup(SourceLookup rootLookup) {
             this.rootLookup.setSource(rootLookup.source());
             this.rootLookup.setSourceContentType(rootLookup.sourceContentType());
         }
+
+        @Override
+        public SearchLookup getSearchLookup() {
+            if (this.searchLookup == null) {
+                SearchExecutionContext sec = getSearchExecutionContext();
+                this.searchLookup = new SearchLookup(sec::getFieldType, (f, l) -> sec.getForField(f));
+            }
+            return this.searchLookup;
+        }
     }
 
     public static void intersect(Weight weight, Weight innerHitQueryWeight, Collector collector, LeafReaderContext ctx) throws IOException {

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

@@ -35,7 +35,7 @@ public InnerHitsPhase(FetchPhase fetchPhase) {
 
     @Override
     public FetchSubPhaseProcessor getProcessor(FetchContext searchContext) {
-        if (searchContext.innerHits() == null) {
+        if (searchContext.innerHits() == null || searchContext.innerHits().getInnerHits().isEmpty()) {
             return null;
         }
         Map<String, InnerHitsContext.InnerHitSubContext> innerHits = searchContext.innerHits().getInnerHits();

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

@@ -36,6 +36,7 @@
 import org.elasticsearch.search.fetch.subphase.InnerHitsContext;
 import org.elasticsearch.search.fetch.subphase.ScriptFieldsContext;
 import org.elasticsearch.search.fetch.subphase.highlight.SearchHighlightContext;
+import org.elasticsearch.search.lookup.SearchLookup;
 import org.elasticsearch.search.profile.Profilers;
 import org.elasticsearch.search.query.QuerySearchResult;
 import org.elasticsearch.search.rescore.RescoreContext;
@@ -335,6 +336,10 @@ public final boolean hasOnlySuggest() {
 
     public abstract SearchExecutionContext getSearchExecutionContext();
 
+    public SearchLookup getSearchLookup() {
+        return getSearchExecutionContext().lookup();
+    }
+
     @Override
     public String toString() {
         StringBuilder result = new StringBuilder().append(shardTarget());

server/src/main/java/org/elasticsearch/search/lookup/ValuesLookup.java

@@ -9,15 +9,35 @@
 package org.elasticsearch.search.lookup;
 
 import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.common.xcontent.XContentType;
 
 import java.util.Map;
+import java.util.Objects;
 
 public interface ValuesLookup {
 
     SourceLookup source();
 
     LeafDocLookup doc();
 
+    static ValuesLookup wrapWithSource(ValuesLookup in, Map<String, Object> source, XContentType xContentType) {
+        System.out.println("Building independent source: " + source);
+        SourceLookup sourceLookup = new SourceLookup();
+        sourceLookup.setSource(source);
+        sourceLookup.setSourceContentType(Objects.requireNonNull(xContentType));
+        return new ValuesLookup() {
+            @Override
+            public SourceLookup source() {
+                return sourceLookup;
+            }
+
+            @Override
+            public LeafDocLookup doc() {
+                return in.doc();
+            }
+        };
+    }
+
     static ValuesLookup sourceOnly(BytesReference source) {
         SourceLookup sourceLookup = new SourceLookup();
         sourceLookup.setSource(source);