highlighter now takes leafReaderContext (#550)

src/main/java/com/yelp/nrtsearch/server/luceneserver/highlights/HighlightFetchTask.java

@@ -29,8 +29,6 @@
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.concurrent.atomic.DoubleAdder;
-import org.apache.lucene.facet.taxonomy.SearcherTaxonomyManager.SearcherAndTaxonomy;
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.Query;
 
@@ -42,18 +40,15 @@ public class HighlightFetchTask implements FetchTask {
 
   private static final double TEN_TO_THE_POWER_SIX = Math.pow(10, 6);
   private final IndexState indexState;
-  private final IndexReader indexReader;
   private final Map<String, HighlightSettings> fieldSettings;
   private final DoubleAdder timeTakenMs = new DoubleAdder();
 
   public HighlightFetchTask(
       IndexState indexState,
-      SearcherAndTaxonomy searcherAndTaxonomy,
       Query searchQuery,
       HighlighterService highlighterService,
       Highlight highlight) {
     this.indexState = indexState;
-    this.indexReader = searcherAndTaxonomy.searcher.getIndexReader();
     this.fieldSettings =
         createPerFieldSettings(highlight, searchQuery, indexState, highlighterService);
     verifyHighlights();
@@ -83,10 +78,10 @@ public void processHit(SearchContext searchContext, LeafReaderContext hitLeaf, B
           (TextBaseFieldDef) fieldDef; // This is safe as we verified earlier
       String[] highlights =
           highlighter.getHighlights(
-              indexReader,
+              hitLeaf,
               fieldSetting.getValue(),
               textBaseFieldDef,
-              hit.getLuceneDocId(),
+              hit.getLuceneDocId() - hitLeaf.docBase,
               searchContext);
       if (highlights != null && highlights.length > 0 && highlights[0] != null) {
         Highlights.Builder builder = Highlights.newBuilder();

src/main/java/com/yelp/nrtsearch/server/luceneserver/highlights/Highlighter.java

@@ -18,7 +18,7 @@
 import com.yelp.nrtsearch.server.luceneserver.field.TextBaseFieldDef;
 import com.yelp.nrtsearch.server.luceneserver.search.SearchContext;
 import java.io.IOException;
-import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.LeafReaderContext;
 
 /**
  * This is the highlighter interface to provide the new highlighters with different implementations.
@@ -37,24 +37,25 @@ public interface Highlighter {
   String getName();
 
   /**
-   * Get highlighted segments for the given docId and field.
+   * Get highlighted segments for the given leafDocId and field.
    *
-   * <p>The indexReader gives the full readabity access to the highlighter. And using the docId and
-   * the fieldName derived from the textBasedFieldDef, the target field can be retrieved.</>
+   * <p>The indexReader gives the full readabity access to the highlighter. And using the leafDocId
+   * and the fieldName derived from the textBasedFieldDef, the target field can be retrieved.</>
    *
-   * @param indexReader the indexerReader has the random access to the entire index
+   * @param hitLeaf the hitLeaf context has the random access to the segment
    * @param settings the highlight settings derived from the search request for this field
    * @param textBaseFieldDef the target's field information
-   * @param docId the target's identifier to retrieve the highlighting document from the indexReader
+   * @param leafDocId the target's identifier to retrieve the highlighting document from the
+   *     leafReaderContext base (luceneId - hitLeaf.docBase)
    * @param searchContext the searchContext to keep the contexts for this search request
    * @return an array of Strings containing all highlighted fragments
    * @throws IOException will be thrown when fail during the document reading
    */
   String[] getHighlights(
-      IndexReader indexReader,
+      LeafReaderContext hitLeaf,
       HighlightSettings settings,
       TextBaseFieldDef textBaseFieldDef,
-      int docId,
+      int leafDocId,
       SearchContext searchContext)
       throws IOException;
 
@@ -68,5 +69,5 @@ String[] getHighlights(
    * @param fieldDef a {@link TextBaseFieldDef} object of the field that is intended to be
    *     highlighted
    */
-  default void verifyFieldIsSupported(TextBaseFieldDef fieldDef) {};
+  default void verifyFieldIsSupported(TextBaseFieldDef fieldDef) {}
 }

src/main/java/com/yelp/nrtsearch/server/luceneserver/highlights/NRTFastVectorHighlighter.java

@@ -20,6 +20,7 @@
 import java.io.IOException;
 import org.apache.lucene.document.FieldType;
 import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.highlight.DefaultEncoder;
 import org.apache.lucene.search.vectorhighlight.BaseFragmentsBuilder;
@@ -64,22 +65,23 @@ public String getName() {
    * Use {@link org.apache.lucene.search.vectorhighlight.FastVectorHighlighter} instance to obtain
    * highlighted fragments for a document.
    *
-   * @param indexReader {@link IndexReader} for the index
+   * @param hitLeaf {@link LeafReaderContext} for the index
    * @param settings {@link HighlightSettings} created from the search request
    * @param textBaseFieldDef Field in document to highlight
-   * @param docId Lucene document ID of the document to highlight
+   * @param leafDocId Lucene document ID of the document to highlight
    * @param _searchContext not in used in fvh
    * @return Array of highlight fragments
    * @throws IOException if there is a low-level IO error
    */
   @Override
   public String[] getHighlights(
-      IndexReader indexReader,
+      LeafReaderContext hitLeaf,
       HighlightSettings settings,
       TextBaseFieldDef textBaseFieldDef,
-      int docId,
+      int leafDocId,
       SearchContext _searchContext)
       throws IOException {
+
     FragListBuilder fragListBuilder;
     int numberOfFragments = settings.getMaxNumFragments();
     int fragmentCharSize = settings.getFragmentSize();
@@ -102,9 +104,9 @@ public String[] getHighlights(
     fragmentsBuilder.setDiscreteMultiValueHighlighting(settings.getDiscreteMultivalue());
 
     return FAST_VECTOR_HIGHLIGHTER.getBestFragments(
-        getFieldQuery(indexReader, settings.getHighlightQuery(), settings.getFieldMatch()),
-        indexReader,
-        docId,
+        getFieldQuery(hitLeaf.reader(), settings.getHighlightQuery(), settings.getFieldMatch()),
+        hitLeaf.reader(),
+        leafDocId,
         textBaseFieldDef.getName(),
         fragmentCharSize,
         numberOfFragments,

src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchRequestProcessor.java

@@ -159,8 +159,7 @@ public static SearchContext buildContextForRequest(
     Highlight highlight = searchRequest.getHighlight();
     if (!highlight.getFieldsList().isEmpty()) {
       HighlightFetchTask highlightFetchTask =
-          new HighlightFetchTask(
-              indexState, searcherAndTaxonomy, query, HighlighterService.getInstance(), highlight);
+          new HighlightFetchTask(indexState, query, HighlighterService.getInstance(), highlight);
       contextBuilder.setHighlightFetchTask(highlightFetchTask);
     }
     contextBuilder.setExtraContext(new ConcurrentHashMap<>());