From 4a1a09cf430bf7c4b5f8d85ce624a9d6df26092e Mon Sep 17 00:00:00 2001
From: Jim Ferenczi <jim.ferenczi@elastic.co>
Date: Fri, 9 Sep 2016 13:00:41 +0200
Subject: [PATCH] Fix highlighting of MultiTermQuery within a
 FunctionScoreQuery

Since the sub query of a function score query is checked on CustomQueryScorer#extractUnknwonQuery we try to extract the terms from the rewritten form of the sub query.
MultiTermQuery rewrites query within a constant score query/weight which returns an empty array when extractTerms is called.
The extraction of the inner terms of a constant score query/weight changed in Lucene somewhere between ES version 2.3 and 2.4 (https://issues.apache.org/jira/browse/LUCENE-6425) which is why this problem occurs on ES > 2.3.
This change moves the extraction of the sub query from CustomQueryScorer#extractUnknownQuery to CustomQueryScorer#extract in order to do the extraction of the terms on the original form of the sub query.
This fixes highlighting of sub queries that extend MultiTermQuery since there is a special path for this kind of query in the QueryScorer (which extract the terms to highlight).
---
 .../subphase/highlight/CustomQueryScorer.java  | 11 +++++------
 .../highlight/HighlighterSearchIT.java         | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/core/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/CustomQueryScorer.java b/core/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/CustomQueryScorer.java
index b62d28f8ab4..4816cba56df 100644
--- a/core/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/CustomQueryScorer.java
+++ b/core/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/CustomQueryScorer.java
@@ -78,10 +78,7 @@ public final class CustomQueryScorer extends QueryScorer {
         @Override
         protected void extractUnknownQuery(Query query,
                                            Map<String, WeightedSpanTerm> terms) throws IOException {
-            if (query instanceof FunctionScoreQuery) {
-                query = ((FunctionScoreQuery) query).getSubQuery();
-                extract(query, 1F, terms);
-            } else if (query instanceof FiltersFunctionScoreQuery) {
+            if (query instanceof FiltersFunctionScoreQuery) {
                 query = ((FiltersFunctionScoreQuery) query).getSubQuery();
                 extract(query, 1F, terms);
             } else if (terms.isEmpty()) {
@@ -97,9 +94,11 @@ public final class CustomQueryScorer extends QueryScorer {
             } else if (query instanceof HasChildQueryBuilder.LateParsingQuery) {
                 // skip has_child or has_parent queries, see: https://github.com/elastic/elasticsearch/issues/14999
                 return;
+            } else if (query instanceof FunctionScoreQuery) {
+                super.extract(((FunctionScoreQuery) query).getSubQuery(), boost, terms);
+            } else {
+                super.extract(query, boost, terms);
             }
-
-            super.extract(query, boost, terms);
         }
     }
 }
diff --git a/core/src/test/java/org/elasticsearch/search/fetch/subphase/highlight/HighlighterSearchIT.java b/core/src/test/java/org/elasticsearch/search/fetch/subphase/highlight/HighlighterSearchIT.java
index 843ab09b2fe..c0fe9bab306 100644
--- a/core/src/test/java/org/elasticsearch/search/fetch/subphase/highlight/HighlighterSearchIT.java
+++ b/core/src/test/java/org/elasticsearch/search/fetch/subphase/highlight/HighlighterSearchIT.java
@@ -38,6 +38,7 @@ import org.elasticsearch.index.query.MultiMatchQueryBuilder;
 import org.elasticsearch.index.query.Operator;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryBuilders;
+import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder;
 import org.elasticsearch.index.search.MatchQuery;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.rest.RestStatus;
@@ -2851,4 +2852,21 @@ public class HighlighterSearchIT extends ESIntegTestCase {
         assertThat(field.getFragments()[0].string(), equalTo("<em>brown</em>"));
         assertThat(field.getFragments()[1].string(), equalTo("<em>cow</em>"));
     }
+
+    public void testFunctionScoreQueryHighlight() throws Exception {
+        client().prepareIndex("test", "type", "1")
+            .setSource(jsonBuilder().startObject().field("text", "brown").endObject())
+            .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
+            .get();
+
+        SearchResponse searchResponse = client().prepareSearch()
+            .setQuery(new FunctionScoreQueryBuilder(QueryBuilders.prefixQuery("text", "bro")))
+            .highlighter(new HighlightBuilder()
+                .field(new Field("text")))
+            .get();
+        assertHitCount(searchResponse, 1);
+        HighlightField field = searchResponse.getHits().getAt(0).highlightFields().get("text");
+        assertThat(field.getFragments().length, equalTo(1));
+        assertThat(field.getFragments()[0].string(), equalTo("<em>brown</em>"));
+    }
 }