Fix Unified highlighter for nested fields when using matchPhrasePrefixQuery

Signed-off-by: sjs004 <[email protected]>

Author: sjs004

CHANGELOG.md

@@ -87,8 +87,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Fix cardinality agg pruning optimization by self collecting ([#19473](https://github.com/opensearch-project/OpenSearch/pull/19473))
 - Implement SslHandler retrieval logic for transport-reactor-netty4 plugin ([#19458](https://github.com/opensearch-project/OpenSearch/pull/19458))
 - Cache serialised cluster state based on cluster state version and node version.([#19307](https://github.com/opensearch-project/OpenSearch/pull/19307))
-
 - Handle negative search request nodes stats ([#19340](https://github.com/opensearch-project/OpenSearch/pull/19340))
+- Fix Unified highlighter for nested fields when using matchPhrasePrefixQuery ([#19442](https://github.com/opensearch-project/OpenSearch/pull/19442))
 
 ### Dependencies
 - Bump `com.gradleup.shadow:shadow-gradle-plugin` from 8.3.5 to 8.3.9 ([#19400](https://github.com/opensearch-project/OpenSearch/pull/19400))

server/src/internalClusterTest/java/org/opensearch/search/fetch/subphase/highlight/HighlighterSearchIT.java

@@ -99,6 +99,7 @@
 import static org.opensearch.index.query.QueryBuilders.constantScoreQuery;
 import static org.opensearch.index.query.QueryBuilders.existsQuery;
 import static org.opensearch.index.query.QueryBuilders.fuzzyQuery;
+import static org.opensearch.index.query.QueryBuilders.matchPhrasePrefixQuery;
 import static org.opensearch.index.query.QueryBuilders.matchPhraseQuery;
 import static org.opensearch.index.query.QueryBuilders.matchQuery;
 import static org.opensearch.index.query.QueryBuilders.multiMatchQuery;
@@ -3518,7 +3519,7 @@ public void testWithNestedQuery() throws Exception {
                 jsonBuilder().startObject()
                     .startArray("foo")
                     .startObject()
-                    .field("text", "brown")
+                    .field("text", "brown cat")
                     .endObject()
                     .startObject()
                     .field("text", "cow")
@@ -3539,7 +3540,7 @@ public void testWithNestedQuery() throws Exception {
             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()[0].string(), equalTo("<em>brown</em> cat"));
             assertThat(field.getFragments()[1].string(), equalTo("<em>cow</em>"));
 
             searchResponse = client().prepareSearch()
@@ -3549,16 +3550,25 @@ public void testWithNestedQuery() throws Exception {
             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>"));
+            assertThat(field.getFragments()[0].string(), equalTo("<em>brown</em> cat"));
 
             searchResponse = client().prepareSearch()
-                .setQuery(nestedQuery("foo", prefixQuery("foo.text", "bro"), ScoreMode.None))
-                .highlighter(new HighlightBuilder().field(new Field("foo.text").highlighterType("plain")))
+                .setQuery(nestedQuery("foo", matchPhraseQuery("foo.text", "brown cat"), 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>"));
+            assertThat(field.getFragments()[0].string(), equalTo("<em>brown</em> <em>cat</em>"));
+
+            searchResponse = client().prepareSearch()
+                .setQuery(nestedQuery("foo", matchPhrasePrefixQuery("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> cat"));
         }
 
         // For unified and fvh highlighters we just check that the nested query is correctly extracted

server/src/main/java/org/opensearch/lucene/search/uhighlight/CustomUnifiedHighlighter.java

@@ -52,6 +52,7 @@
 import org.opensearch.common.Nullable;
 import org.opensearch.common.lucene.search.MultiPhrasePrefixQuery;
 import org.opensearch.index.IndexSettings;
+import org.opensearch.index.search.OpenSearchToParentBlockJoinQuery;
 
 import java.io.IOException;
 import java.text.BreakIterator;
@@ -262,6 +263,8 @@ private Collection<Query> rewriteCustomQuery(Query query) {
             // if original slop is 0 then require inOrder
             boolean inorder = (mpq.getSlop() == 0);
             return Collections.singletonList(new SpanNearQuery(positionSpanQueries, mpq.getSlop() + positionGaps, inorder));
+        } else if (query instanceof OpenSearchToParentBlockJoinQuery) {
+            return Collections.singletonList(((OpenSearchToParentBlockJoinQuery) query).getChildQuery());
         } else {
             return null;
         }