From 0aa6e50e5c9b9c33dff59f1f264de6a38ea17d14 Mon Sep 17 00:00:00 2001
From: Sarthak Nandi <sarthakn@yelp.com>
Date: Tue, 5 Mar 2024 12:05:14 -0800
Subject: [PATCH] Use TotalHitCountCollectorManager instead of
 TopScoreDocCollector if topHits requested <= 0

---
 .../luceneserver/search/SearchCollectorManager.java  |  4 ++--
 .../luceneserver/search/collectors/DocCollector.java |  3 +--
 .../search/collectors/RelevanceCollector.java        | 12 ++++++++----
 3 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchCollectorManager.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchCollectorManager.java
index ac38d3fcd..8dbcb2f40 100644
--- a/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchCollectorManager.java
+++ b/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchCollectorManager.java
@@ -45,7 +45,7 @@ public class SearchCollectorManager
     implements CollectorManager<SearchCollectorManager.SearchCollector, SearcherResult> {
 
   private final DocCollector docCollector;
-  private final CollectorManager<? extends Collector, ? extends TopDocs> docCollectorManger;
+  private final CollectorManager<? extends Collector, ?> docCollectorManger;
   private final List<AdditionalCollectorManager<? extends Collector, ? extends CollectorResult>>
       collectorManagers;
 
@@ -114,7 +114,7 @@ public DocCollector getDocCollector() {
   }
 
   /** Get collector manager implementation used to rank docs */
-  public CollectorManager<? extends Collector, ? extends TopDocs> getDocCollectorManger() {
+  public CollectorManager<? extends Collector, ?> getDocCollectorManger() {
     return docCollectorManger;
   }
 
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..e60a4ab99 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
@@ -34,7 +34,6 @@
 import org.apache.lucene.search.Collector;
 import org.apache.lucene.search.CollectorManager;
 import org.apache.lucene.search.ScoreDoc;
-import org.apache.lucene.search.TopDocs;
 
 /** Abstract base for classes that manage the collection of documents when executing queries. */
 public abstract class DocCollector {
@@ -122,7 +121,7 @@ public boolean terminatedEarly() {
   }
 
   /** Get a lucene level {@link CollectorManager} to rank document for search. */
-  public abstract CollectorManager<? extends Collector, ? extends TopDocs> getManager();
+  public abstract CollectorManager<? extends Collector, ?> getManager();
 
   /**
    * Fill the response hit for the given {@link ScoreDoc}. This method is expected to fill the
diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/collectors/RelevanceCollector.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/collectors/RelevanceCollector.java
index 55284bf5b..419001391 100644
--- a/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/collectors/RelevanceCollector.java
+++ b/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/collectors/RelevanceCollector.java
@@ -24,8 +24,8 @@
 import org.apache.lucene.search.CollectorManager;
 import org.apache.lucene.search.FieldDoc;
 import org.apache.lucene.search.ScoreDoc;
-import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.TopScoreDocCollector;
+import org.apache.lucene.search.TotalHitCountCollectorManager;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -33,7 +33,7 @@
 public class RelevanceCollector extends DocCollector {
   private static final Logger logger = LoggerFactory.getLogger(RelevanceCollector.class);
 
-  private final CollectorManager<TopScoreDocCollector, TopDocs> manager;
+  private final CollectorManager<? extends Collector, ?> manager;
 
   public RelevanceCollector(
       CollectorCreatorContext context,
@@ -52,11 +52,15 @@ public RelevanceCollector(
     } else if (context.getRequest().getTotalHitsThreshold() != 0) {
       totalHitsThreshold = context.getRequest().getTotalHitsThreshold();
     }
-    manager = TopScoreDocCollector.createSharedManager(topHits, searchAfter, totalHitsThreshold);
+    if (topHits <= 0) {
+      manager = new TotalHitCountCollectorManager();
+    } else {
+      manager = TopScoreDocCollector.createSharedManager(topHits, searchAfter, totalHitsThreshold);
+    }
   }
 
   @Override
-  public CollectorManager<? extends Collector, ? extends TopDocs> getManager() {
+  public CollectorManager<? extends Collector, ?> getManager() {
     return manager;
   }