moved originalQueryBuilder to Search instead of Highlight context

Author: piergm

modules/percolator/src/main/java/org/elasticsearch/percolator/PercolatorHighlightSubFetchPhase.java

@@ -87,7 +87,7 @@ public void process(HitContext hit) throws IOException {
                             continue;
                         }
 
-                        SearchHighlightContext highlight = new SearchHighlightContext(fetchContext.highlight().fields(), null);
+                        SearchHighlightContext highlight = new SearchHighlightContext(fetchContext.highlight().fields());
                         FetchSubPhaseProcessor processor = highlightPhase.getProcessor(fetchContext, highlight, query);
                         StoredFieldLoader storedFields = StoredFieldLoader.fromSpec(processor.storedFieldsSpec());
                         LeafStoredFieldLoader leafStoredFields = storedFields.getLoader(percolatorLeafReaderContext, null);

modules/percolator/src/test/java/org/elasticsearch/percolator/PercolatorHighlightSubFetchPhaseTests.java

@@ -45,7 +45,7 @@ public void testHitsExecutionNeeded() {
         );
         PercolatorHighlightSubFetchPhase subFetchPhase = new PercolatorHighlightSubFetchPhase(emptyMap());
         FetchContext fetchContext = mock(FetchContext.class);
-        Mockito.when(fetchContext.highlight()).thenReturn(new SearchHighlightContext(Collections.emptyList(), null));
+        Mockito.when(fetchContext.highlight()).thenReturn(new SearchHighlightContext(Collections.emptyList()));
         Mockito.when(fetchContext.query()).thenReturn(new MatchAllDocsQuery());
 
         assertNull(subFetchPhase.getProcessor(fetchContext));

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

@@ -146,6 +146,7 @@ final class DefaultSearchContext extends SearchContext {
         this.shardTarget = shardTarget;
         this.indexService = readerContext.indexService();
         this.indexShard = readerContext.indexShard();
+        this.originalQueryBuilder = request.source().query();
 
         Engine.Searcher engineSearcher = readerContext.acquireSearcher("search");
         this.searcher = new ContextIndexSearcher(

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

@@ -12,6 +12,7 @@
 import org.elasticsearch.index.mapper.SourceFieldMapper;
 import org.elasticsearch.index.mapper.SourceLoader;
 import org.elasticsearch.index.query.ParsedQuery;
+import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.search.SearchExtBuilder;
 import org.elasticsearch.search.fetch.subphase.FetchDocValuesContext;
@@ -106,6 +107,13 @@ public Query query() {
         return searchContext.query();
     }
 
+    /**
+     * The original query, not rewritten.
+     */
+    public QueryBuilder originalQueryBuilder() {
+        return searchContext.originalQueryBuilder();
+    }
+
     /**
      * The original query in its rewritten form.
      */

server/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/HighlightBuilder.java

@@ -303,7 +303,7 @@ public SearchHighlightContext build(SearchExecutionContext context, QueryBuilder
                 new SearchHighlightContext.Field(field.name(), fieldOptionsBuilder.merge(globalOptionsBuilder.build()).build())
             );
         }
-        return new SearchHighlightContext(fieldOptions, originalQuery);
+        return new SearchHighlightContext(fieldOptions);
     }
 
     /**

server/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/HighlightPhase.java

@@ -141,7 +141,7 @@ private FieldContext contextBuilders(
                     sourceRequired = true;
                 }
 
-                Query highlightQuery = getHighlightQuery(highlightContext, field, fieldType);
+                Query highlightQuery = getHighlightQuery(context, field, fieldType);
 
                 builders.put(
                     fieldName,
@@ -162,12 +162,12 @@ private FieldContext contextBuilders(
     }
 
     private Query getHighlightQuery(
-        SearchHighlightContext highlightContext,
+        FetchContext context,
         SearchHighlightContext.Field field,
         MappedFieldType fieldType
     ) {
         if (fieldType instanceof ConstantFieldType) {
-            QueryBuilder originalQuery = highlightContext.originalQuery();
+            QueryBuilder originalQuery = context.originalQueryBuilder();
             if (originalQuery != null) {
                 return originalQuery.toHighlightQuery(fieldType.name());
             }

server/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/SearchHighlightContext.java

@@ -10,7 +10,6 @@
 
 import org.apache.lucene.search.Query;
 import org.elasticsearch.common.util.Maps;
-import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder.BoundaryScannerType;
 
 import java.util.Arrays;
@@ -22,25 +21,19 @@
 
 public class SearchHighlightContext {
     private final Map<String, Field> fields;
-    private final QueryBuilder originalQuery;
 
-    public SearchHighlightContext(Collection<Field> fields, QueryBuilder originalQuery) {
+    public SearchHighlightContext(Collection<Field> fields) {
         assert fields != null;
         this.fields = Maps.newLinkedHashMapWithExpectedSize(fields.size());
         for (Field field : fields) {
             this.fields.put(field.field, field);
         }
-        this.originalQuery = originalQuery;
     }
 
     public Collection<Field> fields() {
         return fields.values();
     }
 
-    public QueryBuilder originalQuery() {
-        return originalQuery;
-    }
-
     public static class Field {
         private final String field;
         private final FieldOptions fieldOptions;

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

@@ -19,6 +19,7 @@
 import org.elasticsearch.index.cache.bitset.BitsetFilterCache;
 import org.elasticsearch.index.mapper.SourceLoader;
 import org.elasticsearch.index.query.ParsedQuery;
+import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryShardException;
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.index.shard.IndexShard;
@@ -70,6 +71,8 @@ public abstract class SearchContext implements Releasable {
 
     private Query rewriteQuery;
 
+    protected QueryBuilder originalQueryBuilder;
+
     protected SearchContext() {}
 
     public abstract void setTask(SearchShardTask task);
@@ -281,6 +284,10 @@ public Query rewrittenQuery() {
         return rewriteQuery;
     }
 
+    public QueryBuilder originalQueryBuilder() {
+        return originalQueryBuilder;
+    }
+
     public abstract int from();
 
     public abstract SearchContext from(int from);

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

@@ -88,6 +88,7 @@ public TestSearchContext(IndexService indexService) {
         this.indexShard = indexService.getShardOrNull(0);
         searchExecutionContext = indexService.newSearchExecutionContext(0, 0, null, () -> 0L, null, emptyMap());
         this.request = new ShardSearchRequest(indexShard.shardId(), 0L, AliasFilter.EMPTY);
+        this.originalQueryBuilder = request.source().query();
     }
 
     public TestSearchContext(SearchExecutionContext searchExecutionContext) {