Fixed hitsToLog with startHit

src/main/java/com/yelp/nrtsearch/server/handler/SearchHandler.java

@@ -271,16 +271,13 @@ public SearchResponse handle(IndexState indexState, SearchRequest searchRequest)
 
       long t0 = System.nanoTime();
 
-      if (searchContext.getFetchTasks().getHitsLoggerFetchTask() != null) {
-        hits =
-            getHitsFromOffsetAndHitsToLog(
-                hits,
-                searchContext.getStartHit(),
-                searchContext.getTopHits(),
-                searchContext.getHitsToLog());
-      } else {
-        hits = getHitsFromOffset(hits, searchContext.getStartHit(), searchContext.getTopHits());
-      }
+      hits =
+          getHitsFromOffset(
+              hits,
+              searchContext.getStartHit(),
+              Math.max(
+                  searchContext.getTopHits(),
+                  searchContext.getHitsToLog() + searchContext.getStartHit()));
 
       // create Hit.Builder for each hit, and populate with lucene doc id and ranking info
       setResponseHits(searchContext, hits);
@@ -511,12 +508,12 @@ private void fetchFields(SearchContext searchContext)
    *
    * @param hits all hits
    * @param startHit offset into top docs
-   * @param topHits maximum number of hits needed for search response
+   * @param hitsCount maximum number of hits needed for search response
    * @return slice of hits starting at given offset, or empty slice if there are less than startHit
    *     docs
    */
-  public static TopDocs getHitsFromOffset(TopDocs hits, int startHit, int topHits) {
-    int retrieveHits = Math.min(topHits, hits.scoreDocs.length);
+  public static TopDocs getHitsFromOffset(TopDocs hits, int startHit, int hitsCount) {
+    int retrieveHits = Math.min(hitsCount, hits.scoreDocs.length);
     if (startHit != 0 || retrieveHits != hits.scoreDocs.length) {
       // Slice:
       int count = Math.max(0, retrieveHits - startHit);
@@ -529,26 +526,6 @@ public static TopDocs getHitsFromOffset(TopDocs hits, int startHit, int topHits)
     return hits;
   }
 
-  /**
-   * Given all the top documents, produce a slice of the documents starting from a start offset and
-   * going up to the query needed maximum hits. There may be more top docs than the hits limit, if
-   * top docs sampling facets are used. There may be more hits than top hits for logging to be
-   * retrieved.
-   *
-   * @param hits all hits
-   * @param startHit offset into top docs
-   * @param topHits maximum number of hits needed for search response
-   * @param hitsToLog maximum number of hits to be logged
-   * @return slice of hits starting at given offset, or empty slice if there are less than startHit
-   *     docs
-   */
-  public static TopDocs getHitsFromOffsetAndHitsToLog(
-      TopDocs hits, int startHit, int topHits, int hitsToLog) {
-    int retrieveHits = Math.min(Math.max(topHits, hitsToLog), hits.scoreDocs.length);
-    startHit = Math.min(startHit, retrieveHits - hitsToLog);
-    return getHitsFromOffset(hits, startHit, retrieveHits);
-  }
-
   /**
    * Reduce response size by removing any extra hits used for logging. Final search response should
    * only return top hits.

src/main/java/com/yelp/nrtsearch/server/search/collectors/DocCollector.java

@@ -70,7 +70,8 @@ public static int computeNumHitsToCollect(SearchRequest request) {
     // determine how many hits to collect based on request, facets, rescore window and hits to log
     int collectHits = request.getTopHits();
     if (request.hasLoggingHits()) {
-      collectHits = Math.max(collectHits, request.getLoggingHits().getHitsToLog());
+      collectHits =
+          Math.max(collectHits, request.getLoggingHits().getHitsToLog() + request.getStartHit());
     }
     for (Facet facet : request.getFacetsList()) {
       int facetSample = facet.getSampleTopDocs();

src/test/java/com/yelp/nrtsearch/server/logging/HitsLoggerTest.java

@@ -271,7 +271,8 @@ public void testHitsLoggerResponseSizeReductionWithHitsToLogSameAsHitsCount() {
   }
 
   @Test
-  public void testHitsLoggerResponseSizeReductionWithHitsToLogGreaterThanHitsCount() {
+  public void
+      testHitsLoggerResponseSizeReductionWithHitsToLogGreaterThanHitsCountAndLessThanTotalDocs() {
     SearchRequest request =
         SearchRequest.newBuilder()
             .setTopHits(10)
@@ -290,6 +291,33 @@ public void testHitsLoggerResponseSizeReductionWithHitsToLogGreaterThanHitsCount
                 LoggingHits.newBuilder().setName("custom_logger").setHitsToLog(6).build())
             .build();
     SearchResponse response = getGrpcServer().getBlockingStub().search(request);
+    String expectedLogMessage = "LOGGED doc_id: 5, doc_id: 6, doc_id: 7, doc_id: 8, doc_id: 9, ";
+
+    assertEquals(expectedLogMessage, HitsLoggerTest.logMessage);
+    assertEquals(10, response.getTotalHits().getValue());
+    assertEquals(5, response.getHitsCount());
+  }
+
+  @Test
+  public void testHitsLoggerResponseSizeReductionWithHitsToLogGreaterThanHitsCount() {
+    SearchRequest request =
+        SearchRequest.newBuilder()
+            .setTopHits(9)
+            .setStartHit(4)
+            .setIndexName(DEFAULT_TEST_INDEX)
+            .addRetrieveFields("doc_id")
+            .setQuery(
+                Query.newBuilder()
+                    .setTermQuery(
+                        TermQuery.newBuilder()
+                            .setField("vendor_name")
+                            .setTextValue("vendor")
+                            .build())
+                    .build())
+            .setLoggingHits(
+                LoggingHits.newBuilder().setName("custom_logger").setHitsToLog(6).build())
+            .build();
+    SearchResponse response = getGrpcServer().getBlockingStub().search(request);
     String expectedLogMessage =
         "LOGGED doc_id: 4, doc_id: 5, doc_id: 6, doc_id: 7, doc_id: 8, doc_id: 9, ";