Use TotalHitCountCollectorManager instead of TopScoreDocCollector if …

…topHits requested <= 0
Yelp · Mar 5, 2024 · 0aa6e50 · 0aa6e50
1 parent ce67883
commit 0aa6e50
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 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
@@ -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
@@ -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/...ain/java/com/yelp/nrtsearch/server/luceneserver/search/collectors/RelevanceCollector.java b/...ain/java/com/yelp/nrtsearch/server/luceneserver/search/collectors/RelevanceCollector.java
@@ -24,16 +24,16 @@
 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;
 
 /** Collector for getting documents ranked by relevance score. */
 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;
   }