Add unit test & address comments

Signed-off-by: sjs004 <simarjeet.singh004@gmail.com>

Author: sjs004

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

@@ -229,8 +229,7 @@ protected Collection<Query> preSpanQueryRewrite(Query query) {
      * Translate custom queries in queries that are supported by the unified highlighter.
      */
     private Collection<Query> rewriteCustomQuery(Query query) {
-        if (query instanceof MultiPhrasePrefixQuery) {
-            MultiPhrasePrefixQuery mpq = (MultiPhrasePrefixQuery) query;
+        if (query instanceof MultiPhrasePrefixQuery mpq) {
             Term[][] terms = mpq.getTerms();
             int[] positions = mpq.getPositions();
             SpanQuery[] positionSpanQueries = new SpanQuery[positions.length];
@@ -263,13 +262,14 @@ 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 if (query instanceof OpenSearchToParentBlockJoinQuery parentQuery) {
+            return Collections.singletonList(parentQuery.getChildQuery());
         } else {
             return null;
         }
     }
 
+
     /**
      * Forces the offset source for this highlighter
      */

server/src/test/java/org/opensearch/lucene/search/uhighlight/CustomUnifiedHighlighterTests.java

@@ -44,6 +44,7 @@
 import org.apache.lucene.index.IndexOptions;
 import org.apache.lucene.index.IndexWriterConfig;
 import org.apache.lucene.index.Term;
+import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.queries.CommonTermsQuery;
 import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.BooleanQuery;
@@ -56,13 +57,19 @@
 import org.apache.lucene.search.TermQuery;
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.highlight.DefaultEncoder;
+import org.apache.lucene.search.join.BitSetProducer;
+import org.apache.lucene.search.join.ScoreMode;
 import org.apache.lucene.search.uhighlight.UnifiedHighlighter;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.index.RandomIndexWriter;
+import org.apache.lucene.util.BitSet;
+import org.apache.lucene.util.FixedBitSet;
 import org.opensearch.common.lucene.search.MultiPhrasePrefixQuery;
 import org.opensearch.core.common.Strings;
+import org.opensearch.index.search.OpenSearchToParentBlockJoinQuery;
 import org.opensearch.test.OpenSearchTestCase;
 
+import java.io.IOException;
 import java.text.BreakIterator;
 import java.util.Locale;
 
@@ -335,4 +342,34 @@ public void testOverlappingTerms() throws Exception {
         assertHighlightOneDoc("text", inputs, analyzer, query, Locale.ROOT, BreakIterator.getSentenceInstance(Locale.ROOT), 0, outputs);
     }
 
+    public void testOpenSearchToParentBlockJoinQuery() throws Exception {
+        final String[] inputs = { "Nested highlighting query." };
+        final String[] outputs = { "Nested <b>highlighting</b> query." };
+
+        Query childQuery = new TermQuery(new Term("text", "highlighting"));
+        BitSetProducer parentsFilter = new BitSetProducer() {
+            @Override
+            public BitSet getBitSet(LeafReaderContext context) throws IOException {
+                FixedBitSet bits = new FixedBitSet(context.reader().maxDoc());
+                for (int i = 0; i < context.reader().maxDoc(); i++) {
+                    bits.set(i);
+                }
+                return bits;
+            }
+        };
+
+        Query parentQuery = new OpenSearchToParentBlockJoinQuery(childQuery, parentsFilter, ScoreMode.None, "foo");
+        assertHighlightOneDoc(
+            "text",
+            inputs,
+            new StandardAnalyzer(),
+            parentQuery,
+            Locale.ROOT,
+            BreakIterator.getSentenceInstance(Locale.ROOT),
+            0,
+            outputs
+        );
+    }
+
+
 }