Fix nested query highlighting (#26305)

This commit extracts the inner query in the ESToParentBlockJoinQuery for highlighting.
This query has been added in 5.4 and breaks plain highlighting on nested queries.
Highlighters that use postings or term vectors are not affected because they can't highlight nested documents correctly.

Fixes #26230

Author: jimczi

core/src/main/java/org/apache/lucene/search/uhighlight/CustomUnifiedHighlighter.java

@@ -38,6 +38,7 @@
 import org.elasticsearch.common.lucene.all.AllTermQuery;
 import org.elasticsearch.common.lucene.search.MultiPhrasePrefixQuery;
 import org.elasticsearch.common.lucene.search.function.FunctionScoreQuery;
+import org.elasticsearch.index.search.ESToParentBlockJoinQuery;
 
 import java.io.IOException;
 import java.text.BreakIterator;
@@ -210,6 +211,8 @@ private Collection<Query> rewriteCustomQuery(Query query) {
             return Collections.singletonList(new TermQuery(atq.getTerm()));
         } else if (query instanceof FunctionScoreQuery) {
             return Collections.singletonList(((FunctionScoreQuery) query).getSubQuery());
+        } else if (query instanceof ESToParentBlockJoinQuery) {
+            return Collections.singletonList(((ESToParentBlockJoinQuery) query).getChildQuery());
         } else {
             return null;
         }

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

@@ -30,6 +30,7 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.SynonymQuery;
 import org.apache.lucene.search.TermQuery;
+import org.apache.lucene.search.join.ToParentBlockJoinQuery;
 import org.apache.lucene.search.spans.SpanTermQuery;
 import org.elasticsearch.common.lucene.search.MultiPhrasePrefixQuery;
 import org.elasticsearch.common.lucene.search.function.FunctionScoreQuery;
@@ -90,6 +91,11 @@ void flatten(Query sourceQuery, IndexReader reader, Collection<Query> flatQuerie
             for (Term term : synQuery.getTerms()) {
                 flatten(new TermQuery(term), reader, flatQueries, boost);
             }
+        } else if (sourceQuery instanceof ESToParentBlockJoinQuery) {
+            Query childQuery = ((ESToParentBlockJoinQuery) sourceQuery).getChildQuery();
+            if (childQuery != null) {
+                flatten(childQuery, reader, flatQueries, boost);
+            }
         } else {
             super.flatten(sourceQuery, reader, flatQueries, boost);
         }

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

@@ -25,6 +25,7 @@
 import org.apache.lucene.search.highlight.WeightedSpanTerm;
 import org.apache.lucene.search.highlight.WeightedSpanTermExtractor;
 import org.elasticsearch.common.lucene.search.function.FunctionScoreQuery;
+import org.elasticsearch.index.search.ESToParentBlockJoinQuery;
 
 import java.io.IOException;
 import java.util.Map;
@@ -86,6 +87,8 @@ protected void extract(Query query, float boost, Map<String, WeightedSpanTerm> t
                 return;
             } else if (query instanceof FunctionScoreQuery) {
                 super.extract(((FunctionScoreQuery) query).getSubQuery(), boost, terms);
+            } else if (query instanceof ESToParentBlockJoinQuery) {
+                super.extract(((ESToParentBlockJoinQuery) query).getChildQuery(), boost, terms);
             } else {
                 super.extract(query, boost, terms);
             }

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

@@ -20,7 +20,6 @@
 
 import com.carrotsearch.randomizedtesting.generators.RandomPicks;
 import org.apache.lucene.search.join.ScoreMode;
-import org.elasticsearch.Version;
 import org.elasticsearch.action.index.IndexRequestBuilder;
 import org.elasticsearch.action.search.SearchRequestBuilder;
 import org.elasticsearch.action.search.SearchResponse;
@@ -2841,4 +2840,80 @@ public void testHighlightQueryRewriteDatesWithNow() throws Exception {
                 equalTo("<x>hello</x> world"));
         }
     }
+
+    public void testWithNestedQuery() throws Exception {
+        String mapping = jsonBuilder().startObject().startObject("type").startObject("properties")
+            .startObject("text")
+                .field("type", "text")
+                .field("index_options", "offsets")
+                .field("term_vector", "with_positions_offsets")
+            .endObject()
+            .startObject("foo")
+                .field("type", "nested")
+                .startObject("properties")
+                    .startObject("text")
+                        .field("type", "text")
+                    .endObject()
+                .endObject()
+            .endObject()
+            .endObject().endObject().endObject().string();
+        prepareCreate("test").addMapping("type", mapping, XContentType.JSON).get();
+
+        client().prepareIndex("test", "type", "1").setSource(jsonBuilder().startObject()
+            .startArray("foo")
+                .startObject().field("text", "brown").endObject()
+                .startObject().field("text", "cow").endObject()
+            .endArray()
+            .field("text", "brown")
+            .endObject()).setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
+            .get();
+
+        for (String type : new String[] {"unified", "plain"}) {
+            SearchResponse searchResponse = client().prepareSearch()
+                .setQuery(nestedQuery("foo", matchQuery("foo.text", "brown cow"), ScoreMode.None))
+                .highlighter(new HighlightBuilder()
+                    .field(new Field("foo.text").highlighterType(type)))
+                .get();
+            assertHitCount(searchResponse, 1);
+            HighlightField field = searchResponse.getHits().getAt(0).getHighlightFields().get("foo.text");
+            assertThat(field.getFragments().length, equalTo(2));
+            assertThat(field.getFragments()[0].string(), equalTo("<em>brown</em>"));
+            assertThat(field.getFragments()[1].string(), equalTo("<em>cow</em>"));
+
+            searchResponse = client().prepareSearch()
+                .setQuery(nestedQuery("foo", prefixQuery("foo.text", "bro"), ScoreMode.None))
+                .highlighter(new HighlightBuilder()
+                    .field(new Field("foo.text").highlighterType(type)))
+                .get();
+            assertHitCount(searchResponse, 1);
+            field = searchResponse.getHits().getAt(0).getHighlightFields().get("foo.text");
+            assertThat(field.getFragments().length, equalTo(1));
+            assertThat(field.getFragments()[0].string(), equalTo("<em>brown</em>"));
+
+            searchResponse = client().prepareSearch()
+                .setQuery(nestedQuery("foo", prefixQuery("foo.text", "bro"), ScoreMode.None))
+                .highlighter(new HighlightBuilder()
+                    .field(new Field("foo.text").highlighterType("plain")))
+                .get();
+            assertHitCount(searchResponse, 1);
+            field = searchResponse.getHits().getAt(0).getHighlightFields().get("foo.text");
+            assertThat(field.getFragments().length, equalTo(1));
+            assertThat(field.getFragments()[0].string(), equalTo("<em>brown</em>"));
+        }
+
+        // For unified and fvh highlighters we just check that the nested query is correctly extracted
+        // but we highlight the root text field since nested documents cannot be highlighted with postings nor term vectors
+        // directly.
+        for (String type : ALL_TYPES) {
+            SearchResponse searchResponse = client().prepareSearch()
+                .setQuery(nestedQuery("foo", prefixQuery("foo.text", "bro"), ScoreMode.None))
+                .highlighter(new HighlightBuilder()
+                    .field(new Field("text").highlighterType(type).requireFieldMatch(false)))
+                .get();
+            assertHitCount(searchResponse, 1);
+            HighlightField field = searchResponse.getHits().getAt(0).getHighlightFields().get("text");
+            assertThat(field.getFragments().length, equalTo(1));
+            assertThat(field.getFragments()[0].string(), equalTo("<em>brown</em>"));
+        }
+    }
 }