Support MultiPhrasePrefixQuery and MultiPhraseQuery in highlighters

Closes #2596

src/main/java/org/apache/lucene/search/vectorhighlight/CustomFieldQuery.java

@@ -32,6 +32,7 @@
 import java.io.IOException;
 import java.lang.reflect.Field;
 import java.util.Collection;
+import java.util.List;
 
 /**
  *
@@ -86,18 +87,39 @@ void flatten(Query sourceQuery, IndexReader reader, Collection<Query> flatQuerie
             flatten(((FilteredQuery) sourceQuery).getQuery(), reader, flatQueries);
             flatten(((FilteredQuery) sourceQuery).getFilter(), reader, flatQueries);
         } else if (sourceQuery instanceof MultiPhrasePrefixQuery) {
-            try {
-                flatten(sourceQuery.rewrite(reader), reader, flatQueries);
-            } catch (IOException e) {
-                // ignore
-            }
+            flatten(sourceQuery.rewrite(reader), reader, flatQueries);
         } else if (sourceQuery instanceof FiltersFunctionScoreQuery) {
             flatten(((FiltersFunctionScoreQuery) sourceQuery).getSubQuery(), reader, flatQueries);
+        } else if (sourceQuery instanceof MultiPhraseQuery) {
+            MultiPhraseQuery q = ((MultiPhraseQuery) sourceQuery);
+            convertMultiPhraseQuery(0, new int[q.getTermArrays().size()] , q, q.getTermArrays(), q.getPositions(), reader, flatQueries);
         } else {
             super.flatten(sourceQuery, reader, flatQueries);
         }
     }
-
+
+    private void convertMultiPhraseQuery(int currentPos, int[] termsIdx, MultiPhraseQuery orig, List<Term[]> terms, int[] pos, IndexReader reader, Collection<Query> flatQueries) throws IOException {
+        /*
+         * we walk all possible ways and for each path down the MPQ we create a PhraseQuery this is what FieldQuery supports.
+         * It seems expensive but most queries will pretty small.
+         */
+        if (currentPos == terms.size()) {
+            PhraseQuery query = new PhraseQuery();
+            query.setBoost(orig.getBoost());
+            query.setSlop(orig.getSlop());
+            for (int i = 0; i < termsIdx.length; i++) {
+                query.add(terms.get(i)[termsIdx[i]], pos[i]);
+            }
+            this.flatten(query, reader, flatQueries);
+        } else {
+            Term[] t = terms.get(currentPos);
+            for (int i = 0; i < t.length; i++) {
+                termsIdx[currentPos] = i;
+                convertMultiPhraseQuery(currentPos+1, termsIdx, orig, terms, pos, reader, flatQueries);
+            }
+        }
+    }
+
     void flatten(Filter sourceFilter, IndexReader reader, Collection<Query> flatQueries) throws IOException {
         Boolean highlight = highlightFilters.get();
         if (highlight == null || highlight.equals(Boolean.FALSE)) {

src/main/java/org/elasticsearch/common/lucene/search/MultiPhrasePrefixQuery.java

@@ -262,4 +262,8 @@ private boolean termArraysEquals(List<Term[]> termArrays1, List<Term[]> termArra
         }
         return true;
     }
+
+    public String getField() {
+        return field;
+    }
 }

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

@@ -26,6 +26,7 @@
 import org.apache.lucene.search.highlight.QueryScorer;
 import org.apache.lucene.search.highlight.WeightedSpanTerm;
 import org.apache.lucene.search.highlight.WeightedSpanTermExtractor;
+import org.elasticsearch.common.lucene.search.MultiPhrasePrefixQuery;
 import org.elasticsearch.common.lucene.search.function.FiltersFunctionScoreQuery;
 import org.elasticsearch.common.lucene.search.function.FunctionScoreQuery;
 
@@ -93,6 +94,10 @@ protected void extractUnknownQuery(Query query,
             } else if (query instanceof FilteredQuery) {
                 query = ((FilteredQuery) query).getQuery();
                 extract(query, terms);
+            } else if (query instanceof MultiPhrasePrefixQuery) {
+                MultiPhrasePrefixQuery q = ((MultiPhrasePrefixQuery)query);
+                final IndexReader reader = getReaderForField(q.getField());
+                extract(q.rewrite(reader), terms);
             }
         }
 

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

@@ -948,6 +948,48 @@ public void testFSHHighlightAllMvFragments() throws Exception {
         assertThat(response.hits().hits()[0].highlightFields().get("tags").fragments()[1].string(), equalTo("here is another one that is very long and has the <em>tag</em> token near the end"));
     }
 
+    @Test
+    public void testPhrasePrefix() throws ElasticSearchException, IOException {
+        try {
+            client.admin().indices().prepareDelete("test").execute().actionGet();
+        } catch (IndexMissingException e) {
+            // its ok
+        }
+        client.admin().indices().prepareCreate("test").addMapping("type1", type1TermVectorMapping()).execute().actionGet();
+        client.admin().cluster().prepareHealth("test").setWaitForGreenStatus().execute().actionGet();
+
+        client.prepareIndex("test", "type1", "0")
+                .setSource("field0", "The quick brown fox jumps over the lazy dog", "field1", "The quick brown fox jumps over the lazy dog")
+                .setRefresh(true).execute().actionGet();
+        client.prepareIndex("test", "type1", "1")
+        .setSource("field1", "The quick browse button is a fancy thing, right bro?")
+        .setRefresh(true).execute().actionGet();
+        logger.info("--> highlighting and searching on field0");
+        SearchSourceBuilder source = searchSource()
+                .query(matchPhrasePrefixQuery("field0", "quick bro" ))
+                .from(0).size(60).explain(true)
+                .highlight(highlight().field("field0").order("score").preTags("<x>").postTags("</x>"));
+
+        SearchResponse searchResponse = client.search(searchRequest("test").source(source).searchType(QUERY_THEN_FETCH)).actionGet();
+        assertThat("Failures " + Arrays.toString(searchResponse.shardFailures()), searchResponse.shardFailures().length, equalTo(0));
+        assertThat(searchResponse.hits().totalHits(), equalTo(1l));
+
+        assertThat(searchResponse.hits().getAt(0).highlightFields().get("field0").fragments()[0].string(), equalTo("The <x>quick</x> <x>brown</x> fox jumps over the lazy dog"));
+
+        logger.info("--> highlighting and searching on field1");
+        source = searchSource()
+                .query(matchPhrasePrefixQuery("field1","quick bro" ))
+                .from(0).size(60).explain(true)
+                .highlight(highlight().field("field1").order("score").preTags("<x>").postTags("</x>"));
+
+        searchResponse = client.search(searchRequest("test").source(source).searchType(QUERY_THEN_FETCH)).actionGet();
+        assertThat("Failures " + Arrays.toString(searchResponse.shardFailures()), searchResponse.shardFailures().length, equalTo(0));
+        assertThat(searchResponse.hits().totalHits(), equalTo(2l));
+
+        assertThat(searchResponse.hits().getAt(0).highlightFields().get("field1").fragments()[0].string(), equalTo("The <x>quick browse</x> button is a fancy thing, right bro?"));
+        assertThat(searchResponse.hits().getAt(1).highlightFields().get("field1").fragments()[0].string(), equalTo("The <x>quick brown</x> fox jumps over the lazy dog"));
+    }
+
     @Test
     public void testPlainHighlightDifferentFragmenter() throws Exception {
         try {