Convert some complex PriorityQueue implementations to use comparators (#14817)

These require more refactoring than previous PRs. Remaining PriorityQueue subclasses are tricky, are public API, or involve multiple further subclasses

Author: thecoop

lucene/CHANGES.txt

@@ -85,6 +85,8 @@ Other
 * GITHUB#14613: Rewrite APIJAR extractor to use Java 24 classfile API and kill ASM dependency also for build system.  (Uwe Schindler)
 
 * GITHUB#14705: Use Comparators for some PriorityQueue implementations. (Simon Cooper)
+* GITHUB#14761: Use more Comparators for PriorityQueue implementations. (Simon Cooper)
+* GITHUB#14817: Refactor some complex uses of PriorityQueue to use Comparators. (Simon Cooper)
 
 ======================= Lucene 10.3.0 =======================
 

lucene/core/src/java/org/apache/lucene/index/MultiTermsEnum.java

@@ -385,7 +385,7 @@ protected boolean lessThan(TermsEnumWithSlice termsA, TermsEnumWithSlice termsB)
     }
 
     /**
-     * Add the {@link #top()} slice as well as all slices that are positionned on the same term to
+     * Add the {@link #top()} slice as well as all slices that are positioned on the same term to
      * {@code tops} and return how many of them there are.
      */
     int fillTop(TermsEnumWithSlice[] tops) {
@@ -402,7 +402,7 @@ int fillTop(TermsEnumWithSlice[] tops) {
         final int index = stack[--stackLen];
         final int leftChild = index << 1;
         for (int child = leftChild, end = Math.min(size, leftChild + 1); child <= end; ++child) {
-          TermsEnumWithSlice te = get(child);
+          TermsEnumWithSlice te = (TermsEnumWithSlice) getHeapArray()[child];
           if (te.compareTermTo(tops[0]) == 0) {
             tops[numTop++] = te;
             stack[stackLen++] = child;
@@ -411,10 +411,6 @@ int fillTop(TermsEnumWithSlice[] tops) {
       }
       return numTop;
     }
-
-    private TermsEnumWithSlice get(int i) {
-      return (TermsEnumWithSlice) getHeapArray()[i];
-    }
   }
 
   @Override

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

@@ -191,23 +191,20 @@ static MatchesIterator fromSubIterators(List<MatchesIterator> mis) throws IOExce
 
   private DisjunctionMatchesIterator(List<MatchesIterator> matches) throws IOException {
     queue =
-        new PriorityQueue<MatchesIterator>(matches.size()) {
-          @Override
-          protected boolean lessThan(MatchesIterator a, MatchesIterator b) {
-            if (a.startPosition() == -1 && b.startPosition() == -1) {
-              try {
-                return a.startOffset() < b.startOffset()
-                    || (a.startOffset() == b.startOffset() && a.endOffset() < b.endOffset())
-                    || (a.startOffset() == b.startOffset() && a.endOffset() == b.endOffset());
-              } catch (IOException e) {
-                throw new IllegalArgumentException("Failed to retrieve term offset", e);
+        PriorityQueue.usingLessThan(
+            matches.size(),
+            (a, b) -> {
+              if (a.startPosition() == -1 && b.startPosition() == -1) {
+                try {
+                  return a.startOffset() < b.startOffset()
+                      || (a.startOffset() == b.startOffset() && a.endOffset() <= b.endOffset());
+                } catch (IOException e) {
+                  throw new IllegalArgumentException("Failed to retrieve term offset", e);
+                }
               }
-            }
-            return a.startPosition() < b.startPosition()
-                || (a.startPosition() == b.startPosition() && a.endPosition() < b.endPosition())
-                || (a.startPosition() == b.startPosition() && a.endPosition() == b.endPosition());
-          }
-        };
+              return a.startPosition() < b.startPosition()
+                  || (a.startPosition() == b.startPosition() && a.endPosition() <= b.endPosition());
+            });
     for (MatchesIterator mi : matches) {
       if (mi.next()) {
         queue.add(mi);

lucene/core/src/java/org/apache/lucene/util/PriorityQueue.java

@@ -21,6 +21,7 @@
 import java.util.Comparator;
 import java.util.Iterator;
 import java.util.NoSuchElementException;
+import java.util.function.IntFunction;
 import java.util.function.Supplier;
 
 /**
@@ -320,6 +321,30 @@ public final boolean remove(T element) {
     return false;
   }
 
+  /**
+   * Moves the contents of this queue into a new array created by {@code newArray}, lowest items
+   * first
+   */
+  public T[] drainToArrayLowestFirst(IntFunction<T[]> newArray) {
+    T[] array = newArray.apply(size);
+    for (int i = 0; i < array.length; i++) {
+      array[i] = pop();
+    }
+    return array;
+  }
+
+  /**
+   * Moves the contents of this queue into a new array created by {@code newArray}, highest items
+   * first
+   */
+  public T[] drainToArrayHighestFirst(IntFunction<T[]> newArray) {
+    T[] array = newArray.apply(size);
+    for (int i = array.length - 1; i >= 0; i--) {
+      array[i] = pop();
+    }
+    return array;
+  }
+
   private boolean upHeap(int origPos) {
     int i = origPos;
     T node = heap[i]; // save bottom node

lucene/grouping/src/java/org/apache/lucene/search/grouping/BlockGroupingCollector.java

@@ -17,6 +17,7 @@
 package org.apache.lucene.search.grouping;
 
 import java.io.IOException;
+import java.util.Comparator;
 import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.DocIdSetIterator;
@@ -92,7 +93,7 @@ public class BlockGroupingCollector extends SimpleCollector {
   private int groupEndDocID;
   private DocIdSetIterator lastDocPerGroupBits;
   private Scorable scorer;
-  private final GroupQueue groupQueue;
+  private final PriorityQueue<OneGroup> groupQueue;
   private boolean groupCompetes;
 
   private static final class OneGroup {
@@ -105,34 +106,31 @@ private static final class OneGroup {
     int comparatorSlot;
   }
 
-  // Sorts by groupSort.  Not static -- uses comparators, reversed
-  private final class GroupQueue extends PriorityQueue<OneGroup> {
-
-    public GroupQueue(int size) {
-      super(size);
-    }
-
-    @Override
-    protected boolean lessThan(final OneGroup group1, final OneGroup group2) {
-
-      // System.out.println("    ltcheck");
-      assert group1 != group2;
-      assert group1.comparatorSlot != group2.comparatorSlot;
-
-      final int numComparators = comparators.length;
-      for (int compIDX = 0; compIDX < numComparators; compIDX++) {
-        final int c =
-            reversed[compIDX]
-                * comparators[compIDX].compare(group1.comparatorSlot, group2.comparatorSlot);
-        if (c != 0) {
-          // Short circuit
-          return c > 0;
-        }
-      }
-
-      // Break ties by docID; lower docID is always sorted first
-      return group1.topGroupDoc > group2.topGroupDoc;
-    }
+  private PriorityQueue<OneGroup> createGroupQueue(int size) {
+    // Sorts by groupSort
+    return PriorityQueue.usingComparator(
+        size,
+        ((Comparator<OneGroup>)
+                (group1, group2) -> {
+                  assert group1 != group2;
+                  assert group1.comparatorSlot != group2.comparatorSlot;
+
+                  final int numComparators = comparators.length;
+                  for (int compIDX = 0; compIDX < numComparators; compIDX++) {
+                    final int c =
+                        reversed[compIDX]
+                            * comparators[compIDX].compare(
+                                group1.comparatorSlot, group2.comparatorSlot);
+                    if (c != 0) {
+                      // Short circuit
+                      return c;
+                    }
+                  }
+                  return 0;
+                })
+            .thenComparingInt(
+                g -> g.topGroupDoc) // Break ties by docID; lower docID is always sorted first
+            .reversed());
   }
 
   // Called when we transition to another group; if the
@@ -221,7 +219,7 @@ public BlockGroupingCollector(
       throw new IllegalArgumentException("topNGroups must be >= 1 (got " + topNGroups + ")");
     }
 
-    groupQueue = new GroupQueue(topNGroups);
+    groupQueue = createGroupQueue(topNGroups);
     pendingSubDocs = new int[10];
     if (needsScores) {
       pendingSubScores = new float[10];

lucene/highlighter/src/java/org/apache/lucene/search/highlight/Highlighter.java

@@ -292,10 +292,7 @@ public final TextFragment[] getBestTextFragments(
       }
 
       // return the most relevant fragments
-      TextFragment[] frag = new TextFragment[fragQueue.size()];
-      for (int i = frag.length - 1; i >= 0; i--) {
-        frag[i] = fragQueue.pop();
-      }
+      TextFragment[] frag = fragQueue.drainToArrayHighestFirst(TextFragment[]::new);
 
       // merge any contiguous fragments to improve readability
       if (mergeContiguousFragments) {

lucene/highlighter/src/java/org/apache/lucene/search/matchhighlight/PassageSelector.java

@@ -168,10 +168,7 @@ public List<Passage> pickBest(
     // Collect from the priority queue (reverse the order so that highest-scoring are first).
     Passage[] passages;
     if (pq.size() > 0) {
-      passages = new Passage[pq.size()];
-      for (int i = pq.size(); --i >= 0; ) {
-        passages[i] = pq.pop();
-      }
+      passages = pq.drainToArrayHighestFirst(Passage[]::new);
     } else {
       // Handle the default, no highlighting markers case.
       passages = pickDefaultPassage(value, maxPassageWindow, maxPassages, permittedPassageRanges);

lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenByteKnnVectorQuery.java

@@ -148,10 +148,7 @@ protected TopDocs exactSearch(
       queue.pop();
     }
 
-    ScoreDoc[] topScoreDocs = new ScoreDoc[queue.size()];
-    for (int i = topScoreDocs.length - 1; i >= 0; i--) {
-      topScoreDocs[i] = queue.pop();
-    }
+    ScoreDoc[] topScoreDocs = queue.drainToArrayHighestFirst(ScoreDoc[]::new);
 
     TotalHits totalHits = new TotalHits(acceptIterator.cost(), relation);
     return new TopDocs(totalHits, topScoreDocs);

lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenFloatKnnVectorQuery.java

@@ -147,10 +147,7 @@ protected TopDocs exactSearch(
       queue.pop();
     }
 
-    ScoreDoc[] topScoreDocs = new ScoreDoc[queue.size()];
-    for (int i = topScoreDocs.length - 1; i >= 0; i--) {
-      topScoreDocs[i] = queue.pop();
-    }
+    ScoreDoc[] topScoreDocs = queue.drainToArrayHighestFirst(ScoreDoc[]::new);
 
     TotalHits totalHits = new TotalHits(acceptIterator.cost(), relation);
     return new TopDocs(totalHits, topScoreDocs);

lucene/misc/src/java/org/apache/lucene/misc/HighFreqTerms.java

@@ -124,15 +124,7 @@ public static TermStats[] getHighFreqTerms(
       }
     }
 
-    TermStats[] result = new TermStats[tiq.size()];
-    // we want highest first so we read the queue and populate the array
-    // starting at the end and work backwards
-    int count = tiq.size() - 1;
-    while (tiq.size() != 0) {
-      result[count] = tiq.pop();
-      count--;
-    }
-    return result;
+    return tiq.drainToArrayHighestFirst(TermStats[]::new);
   }
 
   /** Compares terms by docTermFreq */