Added support for external query in postings highlighter

It is now possible to highlight an external query using the postings highlighter, relates to #3630

Closes #4121

src/main/java/org/elasticsearch/search/highlight/FastVectorHighlighter.java

@@ -80,13 +80,13 @@ public class FastVectorHighlighter implements Highlighter {
             if (field.requireFieldMatch()) {
                 if (cache.fieldMatchFieldQuery == null) {
                     // we use top level reader to rewrite the query against all readers, with use caching it across hits (and across readers...)
-                    cache.fieldMatchFieldQuery = new CustomFieldQuery(highlighterContext.highlightQuery, hitContext.topLevelReader(), true, field.requireFieldMatch());
+                    cache.fieldMatchFieldQuery = new CustomFieldQuery(highlighterContext.query.originalQuery(), hitContext.topLevelReader(), true, field.requireFieldMatch());
                 }
                 fieldQuery = cache.fieldMatchFieldQuery;
             } else {
                 if (cache.noFieldMatchFieldQuery == null) {
                     // we use top level reader to rewrite the query against all readers, with use caching it across hits (and across readers...)
-                    cache.noFieldMatchFieldQuery = new CustomFieldQuery(highlighterContext.highlightQuery, hitContext.topLevelReader(), true, field.requireFieldMatch());
+                    cache.noFieldMatchFieldQuery = new CustomFieldQuery(highlighterContext.query.originalQuery(), hitContext.topLevelReader(), true, field.requireFieldMatch());
                 }
                 fieldQuery = cache.noFieldMatchFieldQuery;
             }

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

@@ -21,7 +21,6 @@ package org.elasticsearch.search.highlight;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
-import org.apache.lucene.search.Query;
 import org.apache.lucene.index.FieldInfo;
 import org.elasticsearch.ElasticSearchException;
 import org.elasticsearch.ElasticSearchIllegalArgumentException;
@@ -108,11 +107,11 @@ public class HighlightPhase extends AbstractComponent implements FetchSubPhase {
                     throw new ElasticSearchIllegalArgumentException("unknown highlighter type [" + field.highlighterType() + "] for the field [" + fieldName + "]");
                 }
 
-                Query highlightQuery = field.highlightQuery();
-                if (highlightQuery == null) {
-                    // Don't use the context.query() since it might be rewritten, and we need to pass the non rewritten queries to
-                    // let the highlighter handle MultiTerm ones
-                    highlightQuery = context.parsedQuery().query();
+                HighlighterContext.HighlightQuery highlightQuery;
+                if (field.highlightQuery() == null) {
+                    highlightQuery = new HighlighterContext.HighlightQuery(context.parsedQuery().query(), context.query(), context.queryRewritten());
+                } else {
+                    highlightQuery = new HighlighterContext.HighlightQuery(field.highlightQuery(), field.highlightQuery(), false);
                 }
                 HighlighterContext highlighterContext = new HighlighterContext(fieldName, field, fieldMapper, context, hitContext, highlightQuery);
                 HighlightField highlightField = highlighter.highlight(highlighterContext);

src/main/java/org/elasticsearch/search/highlight/HighlighterContext.java

@@ -33,15 +33,39 @@ public class HighlighterContext {
     public final FieldMapper<?> mapper;
     public final SearchContext context;
     public final FetchSubPhase.HitContext hitContext;
-    public final Query highlightQuery;
+    public final HighlightQuery query;
 
     public HighlighterContext(String fieldName, SearchContextHighlight.Field field, FieldMapper<?> mapper, SearchContext context,
-            FetchSubPhase.HitContext hitContext, Query highlightQuery) {
+            FetchSubPhase.HitContext hitContext, HighlightQuery query) {
         this.fieldName = fieldName;
         this.field = field;
         this.mapper = mapper;
         this.context = context;
         this.hitContext = hitContext;
-        this.highlightQuery = highlightQuery;
+        this.query = query;
+    }
+
+    static class HighlightQuery {
+        private final Query originalQuery;
+        private final Query query;
+        private final boolean queryRewritten;
+
+        HighlightQuery(Query originalQuery, Query query, boolean queryRewritten) {
+            this.originalQuery = originalQuery;
+            this.query = query;
+            this.queryRewritten = queryRewritten;
+        }
+
+        public boolean queryRewritten() {
+            return queryRewritten;
+        }
+
+        public Query originalQuery() {
+            return originalQuery;
+        }
+
+        public Query query() {
+            return query;
+        }
     }
 }

src/main/java/org/elasticsearch/search/highlight/PlainHighlighter.java

@@ -68,7 +68,7 @@ public class PlainHighlighter implements Highlighter {
 
         org.apache.lucene.search.highlight.Highlighter entry = cache.get(mapper);
         if (entry == null) {
-            Query query = highlighterContext.highlightQuery;
+            Query query = highlighterContext.query.originalQuery();
             QueryScorer queryScorer = new CustomQueryScorer(query, field.requireFieldMatch() ? mapper.names().indexName() : null);
             queryScorer.setExpandMultiTermQuery(true);
             Fragmenter fragmenter;

src/main/java/org/elasticsearch/search/highlight/PostingsHighlighter.java

@@ -71,7 +71,7 @@ public class PostingsHighlighter implements Highlighter {
             //get the non rewritten query and rewrite it
             Query query;
             try {
-                query = rewrite(context, hitContext.topLevelReader());
+                query = rewrite(highlighterContext, hitContext.topLevelReader());
             } catch (IOException e) {
                 throw new FetchPhaseExecutionException(context, "Failed to highlight field [" + highlighterContext.fieldName + "]", e);
             }
@@ -146,11 +146,11 @@ public class PostingsHighlighter implements Highlighter {
         return null;
     }
 
-    private static Query rewrite(SearchContext searchContext, IndexReader reader) throws IOException {
+    private static Query rewrite(HighlighterContext highlighterContext, IndexReader reader) throws IOException {
         //rewrite is expensive: if the query was already rewritten we try not to rewrite
-        boolean mustRewrite = !searchContext.queryRewritten();
+        boolean mustRewrite = !highlighterContext.query.queryRewritten();
 
-        Query original = searchContext.parsedQuery().query();
+        Query original = highlighterContext.query.originalQuery();
 
         MultiTermQuery originalMultiTermQuery = null;
         MultiTermQuery.RewriteMethod originalRewriteMethod = null;
@@ -166,7 +166,7 @@ public class PostingsHighlighter implements Highlighter {
 
         if (!mustRewrite) {
             //return the rewritten query
-            return searchContext.query();
+            return highlighterContext.query.query();
         }
 
         Query query = original;

src/test/java/org/elasticsearch/search/highlight/HighlighterSearchTests.java

@@ -1301,55 +1301,57 @@ public class HighlighterSearchTests extends AbstractIntegrationTest {
     @Test
     public void testHighlightUsesHighlightQuery() throws IOException {
         assertAcked(prepareCreate("test")
-                .addMapping("type1", "text", "type=string,store=yes,term_vector=with_positions_offsets"));
+                .addMapping("type1", "text", "type=string," + randomStoreField() + "term_vector=with_positions_offsets,index_options=offsets"));
         ensureGreen();
 
-        index("test", "type1", "1", "text", "some stuff stuff stuff stuff stuff to highlight against the stuff phrase");
+        index("test", "type1", "1", "text", "Testing the highlight query feature");
         refresh();
 
-        // Make sure the fvh doesn't highlight in the same way as we're going to do with a scoreQuery because
-        // that would invalidate the test results.
-        Matcher<String> highlightedMatcher = anyOf(
-                containsString("<em>stuff phrase</em>"),            //t FHV normally does this
-                containsString("<em>stuff</em> <em>phrase</em>"));  // Plain normally does this
-        HighlightBuilder.Field field = new HighlightBuilder.Field("text")
-                .fragmentSize(20)
-                .numOfFragments(1)
-                .highlighterType("fvh");
-        SearchRequestBuilder search = client().prepareSearch("test")
-                .setQuery(QueryBuilders.matchQuery("text", "stuff"))
-                .setHighlighterOrder("score")
+        HighlightBuilder.Field field = new HighlightBuilder.Field("text");
+
+        SearchRequestBuilder search = client().prepareSearch("test").setQuery(QueryBuilders.matchQuery("text", "testing"))
                 .addHighlightedField(field);
-        SearchResponse response = search.get();
-        assertHighlight(response, 0, "text", 0, not(highlightedMatcher));
+        Matcher<String> searchQueryMatcher = equalTo("<em>Testing</em> the highlight query feature");
 
-        // And do the same for the plain highlighter
         field.highlighterType("plain");
+        SearchResponse response = search.get();
+        assertHighlight(response, 0, "text", 0, searchQueryMatcher);
+        field.highlighterType("fvh");
+        response = search.get();
+        assertHighlight(response, 0, "text", 0, searchQueryMatcher);
+        field.highlighterType("postings");
         response = search.get();
-        assertHighlight(response, 0, "text", 0, not(highlightedMatcher));
+        assertHighlight(response, 0, "text", 0, searchQueryMatcher);
 
-        // Make sure the fvh takes the highlightQuery into account
-        field.highlighterType("fvh").highlightQuery(matchPhraseQuery("text", "stuff phrase"));
+
+        Matcher<String> hlQueryMatcher = equalTo("Testing the highlight <em>query</em> feature");
+        field.highlightQuery(matchQuery("text", "query"));
+
+        field.highlighterType("fvh");
         response = search.get();
-        assertHighlight(response, 0, "text", 0, highlightedMatcher);
+        assertHighlight(response, 0, "text", 0, hlQueryMatcher);
 
-        // And do the same for the plain highlighter
         field.highlighterType("plain");
         response = search.get();
-        assertHighlight(response, 0, "text", 0, highlightedMatcher);
-        // Note that the plain highlighter doesn't join the highlighted elements for us
+        assertHighlight(response, 0, "text", 0, hlQueryMatcher);
 
-        // Make sure the fvh takes the highlightQuery into account when it is set on the highlight context instead of the field
-        search.setHighlighterQuery(matchPhraseQuery("text", "stuff phrase"));
+        field.highlighterType("postings");
+        response = search.get();
+        assertHighlight(response, 0, "text", 0, hlQueryMatcher);
+
+        // Make sure the the highlightQuery is taken into account when it is set on the highlight context instead of the field
+        search.setHighlighterQuery(matchQuery("text", "query"));
         field.highlighterType("fvh").highlightQuery(null);
         response = search.get();
-        assertHighlight(response, 0, "text", 0, highlightedMatcher);
+        assertHighlight(response, 0, "text", 0, hlQueryMatcher);
 
-        // And do the same for the plain highlighter
         field.highlighterType("plain");
         response = search.get();
-        assertHighlight(response, 0, "text", 0, highlightedMatcher);
-        // Note that the plain highlighter doesn't join the highlighted elements for us
+        assertHighlight(response, 0, "text", 0, hlQueryMatcher);
+
+        field.highlighterType("postings");
+        response = search.get();
+        assertHighlight(response, 0, "text", 0, hlQueryMatcher);
     }
 
     private static String randomStoreField() {