From 1c534fe164557b48ae71beebe1c86c09d6259eaf Mon Sep 17 00:00:00 2001
From: Sarthak Nandi <sarthakn@yelp.com>
Date: Wed, 6 Mar 2024 10:56:38 -0800
Subject: [PATCH] Lucene 9.10 upgrade (#626)

* Bump lucene version to 9.10

* Compute numHitsToCollect before creating Collector
---
 build.gradle                                           |  2 +-
 .../luceneserver/search/SearchRequestProcessor.java    | 10 +++++-----
 .../luceneserver/search/collectors/DocCollector.java   |  6 +++++-
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/build.gradle b/build.gradle
index c33f80c99..651a6330b 100644
--- a/build.gradle
+++ b/build.gradle
@@ -35,7 +35,7 @@ def _artifactId = 'server'
 //This is for https://github.com/gradle/gradle/issues/11308
 System.setProperty("org.gradle.internal.publish.checksums.insecure", "True")
 
-def luceneVersion = '9.9.0'
+def luceneVersion = '9.10.0'
 project.ext.slf4jVersion = '2.0.0-alpha1'
 project.ext.grpcVersion = '1.46.0'
 project.ext.lz4Version = '1.7.0'
diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchRequestProcessor.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchRequestProcessor.java
index ac894d7a4..933a748f7 100644
--- a/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchRequestProcessor.java
+++ b/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchRequestProcessor.java
@@ -387,7 +387,11 @@ private static DocCollector buildDocCollector(CollectorCreatorContext collectorC
                 .collect(Collectors.toList());
 
     DocCollector docCollector;
-    if (searchRequest.getQuery().hasCompletionQuery()) {
+    int numHitsToCollect = DocCollector.computeNumHitsToCollect(searchRequest);
+    // If we don't need hits, just count recalled docs
+    if (numHitsToCollect == 0) {
+      docCollector = new HitCountCollector(collectorCreatorContext, additionalCollectors);
+    } else if (searchRequest.getQuery().hasCompletionQuery()) {
       docCollector = new MyTopSuggestDocsCollector(collectorCreatorContext, additionalCollectors);
     } else if (searchRequest.getQuerySort().getFields().getSortedFieldsList().isEmpty()) {
       if (hasLargeNumHits(searchRequest)) {
@@ -398,10 +402,6 @@ private static DocCollector buildDocCollector(CollectorCreatorContext collectorC
     } else {
       docCollector = new SortFieldCollector(collectorCreatorContext, additionalCollectors);
     }
-    // If we don't need hits, just count recalled docs
-    if (docCollector.getNumHitsToCollect() == 0) {
-      docCollector = new HitCountCollector(collectorCreatorContext, additionalCollectors);
-    }
     return docCollector;
   }
 
diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/collectors/DocCollector.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/collectors/DocCollector.java
index e39bc85e0..1918c3fdc 100644
--- a/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/collectors/DocCollector.java
+++ b/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/collectors/DocCollector.java
@@ -63,6 +63,10 @@ public DocCollector(
     this.indexState = context.getIndexState();
     this.additionalCollectors = additionalCollectors;
 
+    numHitsToCollect = computeNumHitsToCollect(request);
+  }
+
+  public static int computeNumHitsToCollect(SearchRequest request) {
     // determine how many hits to collect based on request, facets and rescore window
     int collectHits = request.getTopHits();
     for (Facet facet : request.getFacetsList()) {
@@ -77,7 +81,7 @@ public DocCollector(
         collectHits = windowSize;
       }
     }
-    numHitsToCollect = collectHits;
+    return collectHits;
   }
 
   /**