skip GeoPointMultiTermQuery when highlighting

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

@@ -27,14 +27,26 @@
 import org.apache.lucene.spatial.geopoint.search.GeoPointInBBoxQuery;
 import org.elasticsearch.common.lucene.search.function.FiltersFunctionScoreQuery;
 import org.elasticsearch.common.lucene.search.function.FunctionScoreQuery;
-import org.elasticsearch.index.query.HasChildQueryBuilder;
 import org.elasticsearch.index.query.HasChildQueryParser;
 
 import java.io.IOException;
 import java.util.Map;
 
 public final class CustomQueryScorer extends QueryScorer {
 
+    private static final Class<?> unsupportedGeoQuery;
+
+    static {
+        try {
+            // in extract() we need to check for GeoPointMultiTermQuery and skip extraction for queries that inherit from it.
+            // But GeoPointMultiTermQuerythat is package private in Lucene hence we cannot use an instanceof check. This is why
+            // we use this rather ugly workaround to get a Class and later be able to compare with isAssignableFrom().
+            unsupportedGeoQuery = Class.forName("org.apache.lucene.spatial.geopoint.search.GeoPointMultiTermQuery");
+        } catch (ClassNotFoundException e) {
+            throw new AssertionError(e);
+        }
+    }
+
     public CustomQueryScorer(Query query, IndexReader reader, String field,
                              String defaultField) {
         super(query, reader, field, defaultField);
@@ -91,7 +103,7 @@ protected void extractUnknownQuery(Query query,
         }
 
         protected void extract(Query query, float boost, Map<String, WeightedSpanTerm> terms) throws IOException {
-            if (query instanceof GeoPointInBBoxQuery) {
+            if (query instanceof GeoPointInBBoxQuery || unsupportedGeoQuery.isAssignableFrom(query.getClass())) {
                 // skip all geo queries, see https://issues.apache.org/jira/browse/LUCENE-7293 and
                 // https://github.com/elastic/elasticsearch/issues/17537
                 return;

core/src/test/java/org/elasticsearch/search/highlight/HighlighterSearchIT.java

@@ -2735,6 +2735,45 @@ public void testGeoFieldHighlightingWithDifferentHighlighters() throws IOExcepti
         assertThat(search.getHits().getAt(0).highlightFields().get("text").fragments().length, equalTo(1));
     }
 
+    public void testGeoFieldHighlightingWhenQueryGetsRewritten() throws IOException {
+        // same as above but in this example the query gets rewritten during highlighting
+        // see https://github.com/elastic/elasticsearch/issues/17537#issuecomment-244939633
+        XContentBuilder mappings = jsonBuilder();
+        mappings.startObject();
+        mappings.startObject("jobs")
+            .startObject("_all")
+            .field("enabled", false)
+            .endObject()
+            .startObject("properties")
+            .startObject("loc")
+            .field("type", "geo_point")
+            .endObject()
+            .startObject("jd")
+            .field("type", "string")
+            .endObject()
+            .endObject()
+            .endObject();
+        mappings.endObject();
+        assertAcked(prepareCreate("test")
+            .addMapping("jobs", mappings));
+        ensureYellow();
+
+        client().prepareIndex("test", "jobs", "1")
+            .setSource(jsonBuilder().startObject().field("jd", "some आवश्यकता है- आर्य समाज अनाथालय, 68 सिविल लाइन्स, बरेली को एक पुरूष रस text")
+                .field("loc", "12.934059,77.610741").endObject())
+            .get();
+        refresh();
+
+        QueryBuilder query = QueryBuilders.functionScoreQuery(QueryBuilders.boolQuery().filter(QueryBuilders.geoBoundingBoxQuery("loc")
+            .bottomRight(-23.065941, 113.610741)
+            .topLeft(48.934059, 41.610741)));
+        SearchResponse search = client().prepareSearch().setSource(
+            new SearchSourceBuilder().query(query)
+                .highlight(new HighlightBuilder().highlighterType("plain").field("jd")).buildAsBytes()).get();
+        assertNoFailures(search);
+        assertThat(search.getHits().totalHits(), equalTo(1L));
+    }
+
     public void testACopyFieldWithNestedQuery() throws Exception {
         String mapping = jsonBuilder().startObject().startObject("type").startObject("properties")
                 .startObject("foo")

core/src/test/java/org/elasticsearch/search/highlight/PlainHighlighterTests.java

@@ -67,8 +67,12 @@ public void checkGeoQueryHighlighting(Query geoQuery) throws IOException, Invali
         String fragment = highlighter.getBestFragment(fieldNameAnalyzer.tokenStream("text", "Arbitrary text field which should not cause " +
             "a failure"), "Arbitrary text field which should not cause a failure");
         assertThat(fragment, equalTo("Arbitrary text field which should not cause a <B>failure</B>"));
-        // TODO: This test will fail if we pass in an instance of GeoPointInBBoxQueryImpl too. Should we also find a way to work around that
-        // or can the query not be rewritten before it is passed into the highlighter?
+        Query rewritten = boolQuery.rewrite(null);
+        highlighter =
+            new org.apache.lucene.search.highlight.Highlighter(new CustomQueryScorer(rewritten));
+        fragment = highlighter.getBestFragment(fieldNameAnalyzer.tokenStream("text", "Arbitrary text field which should not cause " +
+            "a failure"), "Arbitrary text field which should not cause a failure");
+        assertThat(fragment, equalTo("Arbitrary text field which should not cause a <B>failure</B>"));
     }
 
     public void testGeoPointInBBoxQueryHighlighting() throws IOException, InvalidTokenOffsetsException {