Reduce specialization in TopScoreDocCollector. (#14038)

The specialization of `SimpleCollector` vs. `PagingCollector` only helps save a
null check, so it's probably not worth the complexity. Benchmarks cannot see a
difference with this change.

lucene/core/src/java/org/apache/lucene/search/TopScoreDocCollector.java

@@ -29,87 +29,19 @@ import org.apache.lucene.index.LeafReaderContext;
  * <p><b>NOTE</b>: The values {@link Float#NaN} and {@link Float#NEGATIVE_INFINITY} are not valid
  * scores. This collector will not properly collect hits with such scores.
  */
-public abstract class TopScoreDocCollector extends TopDocsCollector<ScoreDoc> {
+public class TopScoreDocCollector extends TopDocsCollector<ScoreDoc> {
-
-  /** Scorable leaf collector */
-  public abstract static class ScorerLeafCollector implements LeafCollector {
-
-    protected Scorable scorer;
-
-    @Override
-    public void setScorer(Scorable scorer) throws IOException {
-      this.scorer = scorer;
-    }
-  }
-
-  static class SimpleTopScoreDocCollector extends TopScoreDocCollector {
-
-    SimpleTopScoreDocCollector(
-        int numHits, int totalHitsThreshold, MaxScoreAccumulator minScoreAcc) {
-      super(numHits, totalHitsThreshold, minScoreAcc);
-    }
-
-    @Override
-    public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
-      // reset the minimum competitive score
-      docBase = context.docBase;
-      minCompetitiveScore = 0f;
-
-      return new ScorerLeafCollector() {
-        @Override
-        public void setScorer(Scorable scorer) throws IOException {
-          super.setScorer(scorer);
-          if (minScoreAcc == null) {
-            updateMinCompetitiveScore(scorer);
-          } else {
-            updateGlobalMinCompetitiveScore(scorer);
-          }
-        }
-
-        @Override
-        public void collect(int doc) throws IOException {
-          float score = scorer.score();
-
-          int hitCountSoFar = ++totalHits;
-
-          if (minScoreAcc != null && (hitCountSoFar & minScoreAcc.modInterval) == 0) {
-            updateGlobalMinCompetitiveScore(scorer);
-          }
-
-          if (score <= pqTop.score) {
-            // Note: for queries that match lots of hits, this is the common case: most hits are not
-            // competitive.
-            if (hitCountSoFar == totalHitsThreshold + 1) {
-              // we just reached totalHitsThreshold, we can start setting the min
-              // competitive score now
-              updateMinCompetitiveScore(scorer);
-            }
-            // Since docs are returned in-order (i.e., increasing doc Id), a document
-            // with equal score to pqTop.score cannot compete since HitQueue favors
-            // documents with lower doc Ids. Therefore reject those docs too.
-          } else {
-            collectCompetitiveHit(doc, score);
-          }
-        }
-
-        private void collectCompetitiveHit(int doc, float score) throws IOException {
-          pqTop.doc = doc + docBase;
-          pqTop.score = score;
-          pqTop = pq.updateTop();
-          updateMinCompetitiveScore(scorer);
-        }
-      };
-    }
-  }
-
-  static class PagingTopScoreDocCollector extends TopScoreDocCollector {
 
   private final ScoreDoc after;
+  final int totalHitsThreshold;
+  final MaxScoreAccumulator minScoreAcc;
 
-    PagingTopScoreDocCollector(
+  // prevents instantiation
+  TopScoreDocCollector(
       int numHits, ScoreDoc after, int totalHitsThreshold, MaxScoreAccumulator minScoreAcc) {
-      super(numHits, totalHitsThreshold, minScoreAcc);
+    super(new HitQueue(numHits, true));
     this.after = after;
+    this.totalHitsThreshold = totalHitsThreshold;
+    this.minScoreAcc = minScoreAcc;
   }
 
   @Override
@@ -133,15 +65,35 @@ public abstract class TopScoreDocCollector extends TopDocsCollector<ScoreDoc> {
   }
 
   @Override
-    public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
+  public ScoreMode scoreMode() {
-      docBase = context.docBase;
+    return totalHitsThreshold == Integer.MAX_VALUE ? ScoreMode.COMPLETE : ScoreMode.TOP_SCORES;
-      final int afterDoc = after.doc - context.docBase;
+  }
-      minCompetitiveScore = 0f;
+
+  @Override
+  public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
+    final int docBase = context.docBase;
+    final ScoreDoc after = this.after;
+    final float afterScore;
+    final int afterDoc;
+    if (after == null) {
+      afterScore = Float.POSITIVE_INFINITY;
+      afterDoc = DocIdSetIterator.NO_MORE_DOCS;
+    } else {
+      afterScore = after.score;
+      afterDoc = after.doc - context.docBase;
+    }
+
+    return new LeafCollector() {
+
+      private Scorable scorer;
+      // HitQueue implements getSentinelObject to return a ScoreDoc, so we know
+      // that at this point top() is already initialized.
+      private ScoreDoc pqTop = pq.top();
+      private float minCompetitiveScore;
 
-      return new ScorerLeafCollector() {
       @Override
       public void setScorer(Scorable scorer) throws IOException {
-          super.setScorer(scorer);
+        this.scorer = scorer;
         if (minScoreAcc == null) {
           updateMinCompetitiveScore(scorer);
         } else {
@@ -159,8 +111,7 @@ public abstract class TopScoreDocCollector extends TopDocsCollector<ScoreDoc> {
           updateGlobalMinCompetitiveScore(scorer);
         }
 
-          float afterScore = after.score;
+        if (after != null && (score > afterScore || (score == afterScore && doc <= afterDoc))) {
-          if (score > afterScore || (score == afterScore && doc <= afterDoc)) {
           // hit was collected on a previous page
           if (totalHitsRelation == TotalHits.Relation.EQUAL_TO) {
             // we just reached totalHitsThreshold, we can start setting the min
@@ -193,42 +144,8 @@ public abstract class TopScoreDocCollector extends TopDocsCollector<ScoreDoc> {
         pqTop = pq.updateTop();
         updateMinCompetitiveScore(scorer);
       }
-      };
-    }
-  }
 
-  int docBase;
+      private void updateGlobalMinCompetitiveScore(Scorable scorer) throws IOException {
-  ScoreDoc pqTop;
-  final int totalHitsThreshold;
-  final MaxScoreAccumulator minScoreAcc;
-  float minCompetitiveScore;
-
-  // prevents instantiation
-  TopScoreDocCollector(int numHits, int totalHitsThreshold, MaxScoreAccumulator minScoreAcc) {
-    super(new HitQueue(numHits, true));
-
-    // HitQueue implements getSentinelObject to return a ScoreDoc, so we know
-    // that at this point top() is already initialized.
-    pqTop = pq.top();
-    this.totalHitsThreshold = totalHitsThreshold;
-    this.minScoreAcc = minScoreAcc;
-  }
-
-  @Override
-  protected TopDocs newTopDocs(ScoreDoc[] results, int start) {
-    if (results == null) {
-      return EMPTY_TOPDOCS;
-    }
-
-    return new TopDocs(new TotalHits(totalHits, totalHitsRelation), results);
-  }
-
-  @Override
-  public ScoreMode scoreMode() {
-    return totalHitsThreshold == Integer.MAX_VALUE ? ScoreMode.COMPLETE : ScoreMode.TOP_SCORES;
-  }
-
-  protected void updateGlobalMinCompetitiveScore(Scorable scorer) throws IOException {
         assert minScoreAcc != null;
         long maxMinScore = minScoreAcc.getRaw();
         if (maxMinScore != Long.MIN_VALUE) {
@@ -245,14 +162,12 @@ public abstract class TopScoreDocCollector extends TopDocsCollector<ScoreDoc> {
         }
       }
 
-  protected void updateMinCompetitiveScore(Scorable scorer) throws IOException {
+      private void updateMinCompetitiveScore(Scorable scorer) throws IOException {
         if (totalHits > totalHitsThreshold) {
-      // since we tie-break on doc id and collect in doc id order, we can require
+          // since we tie-break on doc id and collect in doc id order, we can require the next float
-      // the next float
           // pqTop is never null since TopScoreDocCollector fills the priority queue with sentinel
-      // values
+          // values if the top element is a sentinel value, its score will be -Infty and the below
-      // if the top element is a sentinel value, its score will be -Infty and the below logic is
+          // logic is still valid
-      // still valid
           float localMinScore = Math.nextUp(pqTop.score);
           if (localMinScore > minCompetitiveScore) {
             scorer.setMinCompetitiveScore(localMinScore);
@@ -266,4 +181,6 @@ public abstract class TopScoreDocCollector extends TopDocsCollector<ScoreDoc> {
           }
         }
       }
+    };
+  }
 }

lucene/core/src/java/org/apache/lucene/search/TopScoreDocCollectorManager.java

@@ -125,13 +125,7 @@ public class TopScoreDocCollectorManager
 
   @Override
   public TopScoreDocCollector newCollector() {
-    if (after == null) {
+    return new TopScoreDocCollector(numHits, after, totalHitsThreshold, minScoreAcc);
-      return new TopScoreDocCollector.SimpleTopScoreDocCollector(
-          numHits, totalHitsThreshold, minScoreAcc);
-    } else {
-      return new TopScoreDocCollector.PagingTopScoreDocCollector(
-          numHits, after, totalHitsThreshold, minScoreAcc);
-    }
   }
 
   @Override

lucene/sandbox/src/java/org/apache/lucene/sandbox/search/LargeNumHitsTopDocsCollector.java

@@ -32,7 +32,6 @@ import org.apache.lucene.search.Scorable;
 import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.TopDocs;
-import org.apache.lucene.search.TopScoreDocCollector;
 import org.apache.lucene.search.TotalHits;
 
 /**
@@ -63,11 +62,13 @@ public final class LargeNumHitsTopDocsCollector implements Collector {
   @Override
   public LeafCollector getLeafCollector(LeafReaderContext context) {
     final int docBase = context.docBase;
-    return new TopScoreDocCollector.ScorerLeafCollector() {
+    return new LeafCollector() {
+
+      private Scorable scorer;
 
       @Override
       public void setScorer(Scorable scorer) throws IOException {
-        super.setScorer(scorer);
+        this.scorer = scorer;
       }
 
       @Override