Move HitQueue in TopScoreDocCollector to a LongHeap

lucene/CHANGES.txt

@@ -128,6 +128,8 @@ Optimizations
 
 * GITHUB#14674: Optimize AbstractKnnVectorQuery#createBitSet with intoBitset. (Guo Feng)
 
+* GITHUB#14714: Move HitQueue in TopScoreDocCollector to a LongHeap. (Guo Feng)
+
 Bug Fixes
 ---------------------
 * GITHUB#14654: ValueSource.fromDoubleValuesSource(dvs).getSortField() would throw errors when

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

@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.search;
+
+import org.apache.lucene.util.NumericUtils;
+
+/**
+ * An encoder do encode (doc, score) pair as a long whose sort order is same as {@code (o1, o2) ->
+ * Float.compare(o1.score, o2.score)).thenComparing(Comparator.comparingInt((ScoreDoc o) ->
+ * o.doc).reversed())}
+ *
+ * <p>Note that negative score is allowed but relationship between two codes encoded by negative
+ * scores is undefined. The only thing guaranteed is codes encoded from negative scores are smaller
+ * than codes encoded from non-negative scores.
+ */
+class DocScoreEncoder {
+
+  static final long LEAST_COMPETITIVE_CODE = encode(Integer.MAX_VALUE, Float.NEGATIVE_INFINITY);
+  private static final int POS_INF_TO_SORTABLE_INT = scoreToSortableInt(Float.POSITIVE_INFINITY);
+
+  static long encode(int docId, float score) {
+    return encodeIntScore(docId, scoreToSortableInt(score));
+  }
+
+  static long encodeIntScore(int docId, int score) {
+    return (((long) score) << 32) | (~docId & 0xFFFFFFFFL);
+  }
+
+  static float toScore(long value) {
+    return sortableIntToScore(toIntScore(value));
+  }
+
+  static int toIntScore(long value) {
+    return (int) (value >>> 32);
+  }
+
+  static int docId(long value) {
+    return (int) ~value;
+  }
+
+  static int nextUp(int intScore) {
+    assert intScore <= POS_INF_TO_SORTABLE_INT;
+    int nextUp = Math.min(POS_INF_TO_SORTABLE_INT, intScore + 1);
+    assert nextUp == scoreToSortableInt(Math.nextUp(sortableIntToScore(intScore)));
+    return nextUp;
+  }
+
+  /**
+   * Score is non-negative float so wo use floatToRawIntBits instead of {@link
+   * NumericUtils#floatToSortableInt}. We do not assert score >= 0 here to allow pass negative float
+   * to indicate totally non-competitive, e.g. {@link #LEAST_COMPETITIVE_CODE}.
+   */
+  static int scoreToSortableInt(float score) {
+    assert Float.isNaN(score) == false;
+    return Float.floatToRawIntBits(score);
+  }
+
+  /**
+   * @see #scoreToSortableInt(float)
+   */
+  static float sortableIntToScore(int scoreBits) {
+    float score = Float.intBitsToFloat(scoreBits);
+    assert Float.isNaN(score) == false;
+    return score;
+  }
+}

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

@@ -22,10 +22,10 @@
 /** Maintains the maximum score and its corresponding document id concurrently */
 final class MaxScoreAccumulator {
   // we use 2^10-1 to check the remainder with a bitwise operation
-  static final int DEFAULT_INTERVAL = 0x3ff;
+  private static final int DEFAULT_INTERVAL = 0x3ff;
 
   // scores are always positive
-  final LongAccumulator acc = new LongAccumulator(MaxScoreAccumulator::maxEncode, Long.MIN_VALUE);
+  final LongAccumulator acc = new LongAccumulator(Math::max, Long.MIN_VALUE);
 
   // non-final and visible for tests
   long modInterval;
@@ -34,35 +34,8 @@ final class MaxScoreAccumulator {
     this.modInterval = DEFAULT_INTERVAL;
   }
 
-  /**
-   * Return the max encoded docId and score found in the two longs, following the encoding in {@link
-   * #accumulate}.
-   */
-  private static long maxEncode(long v1, long v2) {
-    float score1 = Float.intBitsToFloat((int) (v1 >> 32));
-    float score2 = Float.intBitsToFloat((int) (v2 >> 32));
-    int cmp = Float.compare(score1, score2);
-    if (cmp == 0) {
-      // tie-break on the minimum doc base
-      return (int) v1 < (int) v2 ? v1 : v2;
-    } else if (cmp > 0) {
-      return v1;
-    }
-    return v2;
-  }
-
-  void accumulate(int docId, float score) {
-    assert docId >= 0 && score >= 0;
-    long encode = (((long) Float.floatToIntBits(score)) << 32) | docId;
-    acc.accumulate(encode);
-  }
-
-  public static float toScore(long value) {
-    return Float.intBitsToFloat((int) (value >> 32));
-  }
-
-  public static int docId(long value) {
-    return (int) value;
+  void accumulate(long code) {
+    acc.accumulate(code);
   }
 
   long getRaw() {

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

@@ -153,18 +153,25 @@ public TopDocs topDocs(int start, int howMany) {
     howMany = Math.min(size - start, howMany);
     ScoreDoc[] results = new ScoreDoc[howMany];
 
-    // pq's pop() returns the 'least' element in the queue, therefore need
-    // to discard the first ones, until we reach the requested range.
+    // Prune the least competitive hits until we reach the requested range.
     // Note that this loop will usually not be executed, since the common usage
     // should be that the caller asks for the last howMany results. However it's
     // needed here for completeness.
-    for (int i = pq.size() - start - howMany; i > 0; i--) {
-      pq.pop();
-    }
+    pruneLeastCompetitiveHitsTo(start + howMany);
 
     // Get the requested results from pq.
     populateResults(results, howMany);
 
     return newTopDocs(results, start);
   }
+
+  /**
+   * Prune the least competitive hits until the number of candidates is less than or equal to {@code
+   * keep}. This is typically called before {@link #populateResults} to ensure we are at right pos.
+   */
+  protected void pruneLeastCompetitiveHitsTo(int keep) {
+    for (int i = pq.size() - keep; i > 0; i--) {
+      pq.pop();
+    }
+  }
 }

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

@@ -367,7 +367,7 @@ protected void updateGlobalMinCompetitiveScore(Scorable scorer) throws IOExcepti
       long maxMinScore = minScoreAcc.getRaw();
       float score;
       if (maxMinScore != Long.MIN_VALUE
-          && (score = MaxScoreAccumulator.toScore(maxMinScore)) > minCompetitiveScore) {
+          && (score = DocScoreEncoder.toScore(maxMinScore)) > minCompetitiveScore) {
         scorer.setMinCompetitiveScore(score);
         minCompetitiveScore = score;
         totalHitsRelation = TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO;
@@ -384,7 +384,7 @@ protected void updateMinCompetitiveScore(Scorable scorer) throws IOException {
         minCompetitiveScore = minScore;
         totalHitsRelation = TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO;
         if (minScoreAcc != null) {
-          minScoreAcc.accumulate(docBase, minScore);
+          minScoreAcc.accumulate(DocScoreEncoder.encode(docBase, minScore));
         }
       }
     }

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

@@ -17,7 +17,9 @@
 package org.apache.lucene.search;
 
 import java.io.IOException;
+import java.util.stream.IntStream;
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.util.LongHeap;
 
 /**
  * A {@link Collector} implementation that collects the top-scoring hits, returning them as a {@link
@@ -32,31 +34,21 @@
 public class TopScoreDocCollector extends TopDocsCollector<ScoreDoc> {
 
   private final ScoreDoc after;
+  private final LongHeap heap;
   final int totalHitsThreshold;
   final MaxScoreAccumulator minScoreAcc;
 
   // prevents instantiation
   TopScoreDocCollector(
       int numHits, ScoreDoc after, int totalHitsThreshold, MaxScoreAccumulator minScoreAcc) {
-    super(new HitQueue(numHits, true));
+    super(null);
+    this.heap = new LongHeap(numHits);
+    IntStream.range(0, numHits).forEach(_ -> heap.push(DocScoreEncoder.LEAST_COMPETITIVE_CODE));
     this.after = after;
     this.totalHitsThreshold = totalHitsThreshold;
     this.minScoreAcc = minScoreAcc;
   }
 
-  @Override
-  protected int topDocsSize() {
-    // Note: this relies on sentinel values having Integer.MAX_VALUE as a doc ID.
-    int[] validTopHitCount = new int[1];
-    pq.forEach(
-        scoreDoc -> {
-          if (scoreDoc.doc != Integer.MAX_VALUE) {
-            validTopHitCount[0]++;
-          }
-        });
-    return validTopHitCount[0];
-  }
-
   @Override
   protected TopDocs newTopDocs(ScoreDoc[] results, int start) {
     return results == null
@@ -73,23 +65,22 @@ public ScoreMode scoreMode() {
   public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
     final int docBase = context.docBase;
     final ScoreDoc after = this.after;
-    final float afterScore;
+    final int afterScore;
     final int afterDoc;
     if (after == null) {
-      afterScore = Float.POSITIVE_INFINITY;
+      afterScore = Integer.MAX_VALUE;
       afterDoc = DocIdSetIterator.NO_MORE_DOCS;
     } else {
-      afterScore = after.score;
+      afterScore = DocScoreEncoder.scoreToSortableInt(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;
+      private long topCode = heap.top();
+      private int topScore = DocScoreEncoder.toIntScore(topCode);
+      private int minCompetitiveScore;
 
       @Override
       public void setScorer(Scorable scorer) throws IOException {
@@ -103,7 +94,7 @@ public void setScorer(Scorable scorer) throws IOException {
 
       @Override
       public void collect(int doc) throws IOException {
-        float score = scorer.score();
+        final int score = DocScoreEncoder.scoreToSortableInt(scorer.score());
 
         int hitCountSoFar = ++totalHits;
 
@@ -121,7 +112,7 @@ public void collect(int doc) throws IOException {
           return;
         }
 
-        if (score <= pqTop.score) {
+        if (score <= topScore) {
           // Note: for queries that match lots of hits, this is the common case: most hits are not
           // competitive.
           if (hitCountSoFar == totalHitsThreshold + 1) {
@@ -138,10 +129,10 @@ public void collect(int doc) throws IOException {
         }
       }
 
-      private void collectCompetitiveHit(int doc, float score) throws IOException {
-        pqTop.doc = doc + docBase;
-        pqTop.score = score;
-        pqTop = pq.updateTop();
+      private void collectCompetitiveHit(int doc, int score) throws IOException {
+        final long code = DocScoreEncoder.encodeIntScore(doc + docBase, score);
+        topCode = heap.updateTop(code);
+        topScore = DocScoreEncoder.toIntScore(topCode);
         updateMinCompetitiveScore(scorer);
       }
 
@@ -152,10 +143,11 @@ private void updateGlobalMinCompetitiveScore(Scorable scorer) throws IOException
           // since we tie-break on doc id and collect in doc id order we can require
           // the next float if the global minimum score is set on a document id that is
           // smaller than the ids in the current leaf
-          float score = MaxScoreAccumulator.toScore(maxMinScore);
-          score = docBase >= MaxScoreAccumulator.docId(maxMinScore) ? Math.nextUp(score) : score;
+          int score = DocScoreEncoder.toIntScore(maxMinScore);
+          score =
+              docBase >= DocScoreEncoder.docId(maxMinScore) ? DocScoreEncoder.nextUp(score) : score;
           if (score > minCompetitiveScore) {
-            scorer.setMinCompetitiveScore(score);
+            scorer.setMinCompetitiveScore(DocScoreEncoder.sortableIntToScore(score));
             minCompetitiveScore = score;
             totalHitsRelation = TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO;
           }
@@ -164,23 +156,44 @@ private void updateGlobalMinCompetitiveScore(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 the next float
-          // pqTop is never null since TopScoreDocCollector fills the priority queue with sentinel
-          // values if the top element is a sentinel value, its score will be -Infty and the below
-          // logic is still valid
-          float localMinScore = Math.nextUp(pqTop.score);
-          if (localMinScore > minCompetitiveScore) {
-            scorer.setMinCompetitiveScore(localMinScore);
+          if (topScore >= minCompetitiveScore) {
+            minCompetitiveScore = DocScoreEncoder.nextUp(topScore);
+            scorer.setMinCompetitiveScore(DocScoreEncoder.sortableIntToScore(minCompetitiveScore));
             totalHitsRelation = TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO;
-            minCompetitiveScore = localMinScore;
             if (minScoreAcc != null) {
               // we don't use the next float but we register the document id so that other leaves or
               // leaf partitions can require it if they are after the current maximum
-              minScoreAcc.accumulate(pqTop.doc, pqTop.score);
+              minScoreAcc.accumulate(topCode);
             }
           }
         }
       }
     };
   }
+
+  @Override
+  protected int topDocsSize() {
+    int cnt = 0;
+    for (int i = 1; i <= heap.size(); i++) {
+      if (heap.get(i) != DocScoreEncoder.LEAST_COMPETITIVE_CODE) {
+        cnt++;
+      }
+    }
+    return cnt;
+  }
+
+  @Override
+  protected void populateResults(ScoreDoc[] results, int howMany) {
+    for (int i = howMany - 1; i >= 0; i--) {
+      long encode = heap.pop();
+      results[i] = new ScoreDoc(DocScoreEncoder.docId(encode), DocScoreEncoder.toScore(encode));
+    }
+  }
+
+  @Override
+  protected void pruneLeastCompetitiveHitsTo(int keep) {
+    for (int i = heap.size() - keep; i > 0; i--) {
+      heap.pop();
+    }
+  }
 }