Swap out some simple PriorityQueue subclasses for one using a Comparator (#14705)

Author: thecoop

lucene/CHANGES.txt

@@ -70,6 +70,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)
+
 ======================= Lucene 10.3.0 =======================
 
 API Changes

lucene/benchmark/src/java/org/apache/lucene/benchmark/quality/utils/QualityQueriesFinder.java

@@ -18,6 +18,7 @@
 
 import java.io.IOException;
 import java.nio.file.Paths;
+import java.util.Comparator;
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.MultiTerms;
@@ -99,7 +100,8 @@ private static String formatQueryAsTrecTopic(
   }
 
   private String[] bestTerms(String field, int numTerms) throws IOException {
-    PriorityQueue<TermDf> pq = new TermsDfQueue(numTerms);
+    PriorityQueue<TermDf> pq =
+        PriorityQueue.usingComparator(numTerms, Comparator.comparingInt(tdf -> tdf.df));
     IndexReader ir = DirectoryReader.open(dir);
     try {
       int threshold = ir.maxDoc() / 10; // ignore words too common.
@@ -136,15 +138,4 @@ private static class TermDf {
       this.df = freq;
     }
   }
-
-  private static class TermsDfQueue extends PriorityQueue<TermDf> {
-    TermsDfQueue(int maxSize) {
-      super(maxSize);
-    }
-
-    @Override
-    protected boolean lessThan(TermDf tf1, TermDf tf2) {
-      return tf1.df < tf2.df;
-    }
-  }
 }

lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90CompoundFormat.java

@@ -17,6 +17,9 @@
 package org.apache.lucene.codecs.lucene90;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.List;
 import org.apache.lucene.codecs.CodecUtil;
 import org.apache.lucene.codecs.CompoundDirectory;
 import org.apache.lucene.codecs.CompoundFormat;
@@ -27,7 +30,6 @@
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexOutput;
-import org.apache.lucene.util.PriorityQueue;
 
 /**
  * Lucene 9.0 compound file format
@@ -105,29 +107,18 @@ public void write(Directory dir, SegmentInfo si, IOContext context) throws IOExc
 
   private record SizedFile(String name, long length) {}
 
-  private static class SizedFileQueue extends PriorityQueue<SizedFile> {
-    SizedFileQueue(int maxSize) {
-      super(maxSize);
-    }
-
-    @Override
-    protected boolean lessThan(SizedFile sf1, SizedFile sf2) {
-      return sf1.length < sf2.length;
-    }
-  }
-
   private void writeCompoundFile(
       IndexOutput entries, IndexOutput data, Directory dir, SegmentInfo si) throws IOException {
     // write number of files
     int numFiles = si.files().size();
     entries.writeVInt(numFiles);
     // first put files in ascending size order so small files fit more likely into one page
-    SizedFileQueue pq = new SizedFileQueue(numFiles);
+    List<SizedFile> files = new ArrayList<>(numFiles);
     for (String filename : si.files()) {
-      pq.add(new SizedFile(filename, dir.fileLength(filename)));
+      files.add(new SizedFile(filename, dir.fileLength(filename)));
     }
-    while (pq.size() > 0) {
-      SizedFile sizedFile = pq.pop();
+    files.sort(Comparator.comparingLong(SizedFile::length));
+    for (SizedFile sizedFile : files) {
       String file = sizedFile.name;
       // align file start offset
       long startOffset = data.alignFilePointer(Long.BYTES);

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

@@ -46,18 +46,6 @@ public class OrdinalMap implements Accountable {
   // need it
   // TODO: use more efficient packed ints structures?
 
-  private static class TermsEnumPriorityQueue extends PriorityQueue<TermsEnumIndex> {
-
-    TermsEnumPriorityQueue(int size) {
-      super(size);
-    }
-
-    @Override
-    protected boolean lessThan(TermsEnumIndex a, TermsEnumIndex b) {
-      return a.compareTermTo(b) < 0;
-    }
-  }
-
   private static class SegmentMap implements Accountable {
     private static final long BASE_RAM_BYTES_USED =
         RamUsageEstimator.shallowSizeOfInstance(SegmentMap.class);
@@ -265,7 +253,8 @@ public static OrdinalMap build(
     long[] segmentOrds = new long[subs.length];
 
     // Just merge-sorts by term:
-    TermsEnumPriorityQueue queue = new TermsEnumPriorityQueue(subs.length);
+    PriorityQueue<TermsEnumIndex> queue =
+        PriorityQueue.usingComparator(subs.length, TermsEnumIndex::compareTermTo);
 
     for (int i = 0; i < subs.length; i++) {
       TermsEnumIndex sub = new TermsEnumIndex(subs[segmentMap.newToOld(i)], i);

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

@@ -18,7 +18,7 @@
 
 import java.io.IOException;
 import java.util.Collection;
-import java.util.Objects;
+import java.util.Comparator;
 import org.apache.lucene.internal.hppc.LongArrayList;
 import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.FixedBitSet;
@@ -40,43 +40,14 @@ static class Bucket {
     int freq;
   }
 
-  static final class HeadPriorityQueue extends PriorityQueue<DisiWrapper> {
-
-    public HeadPriorityQueue(int maxSize) {
-      super(maxSize);
-    }
-
-    @Override
-    protected boolean lessThan(DisiWrapper a, DisiWrapper b) {
-      return a.doc < b.doc;
-    }
-  }
-
-  static final class TailPriorityQueue extends PriorityQueue<DisiWrapper> {
-
-    public TailPriorityQueue(int maxSize) {
-      super(maxSize);
-    }
-
-    @Override
-    protected boolean lessThan(DisiWrapper a, DisiWrapper b) {
-      return a.cost < b.cost;
-    }
-
-    public DisiWrapper get(int i) {
-      Objects.checkIndex(i, size());
-      return (DisiWrapper) getHeapArray()[1 + i];
-    }
-  }
-
   // One bucket per doc ID in the window, non-null if scores are needed or if frequencies need to be
   // counted
   final Bucket[] buckets;
   final FixedBitSet matching = new FixedBitSet(SIZE);
 
   final DisiWrapper[] leads;
-  final HeadPriorityQueue head;
-  final TailPriorityQueue tail;
+  final PriorityQueue<DisiWrapper> head;
+  final PriorityQueue<DisiWrapper> tail;
   final Score score = new Score();
   final int minShouldMatch;
   final long cost;
@@ -101,8 +72,11 @@ public DisiWrapper get(int i) {
       buckets = null;
     }
     this.leads = new DisiWrapper[scorers.size()];
-    this.head = new HeadPriorityQueue(scorers.size() - minShouldMatch + 1);
-    this.tail = new TailPriorityQueue(minShouldMatch - 1);
+    this.head =
+        PriorityQueue.usingComparator(
+            scorers.size() - minShouldMatch + 1, Comparator.comparingInt(d -> d.doc));
+    this.tail =
+        PriorityQueue.usingComparator(minShouldMatch - 1, Comparator.comparingLong(d -> d.cost));
     this.minShouldMatch = minShouldMatch;
     this.needsScores = needsScores;
     LongArrayList costs = new LongArrayList(scorers.size());
@@ -204,8 +178,8 @@ private void scoreWindowIntoBitSetAndReplay(
 
   private DisiWrapper advance(int min) throws IOException {
     assert tail.size() == minShouldMatch - 1;
-    final HeadPriorityQueue head = this.head;
-    final TailPriorityQueue tail = this.tail;
+    final PriorityQueue<DisiWrapper> head = this.head;
+    final PriorityQueue<DisiWrapper> tail = this.tail;
     DisiWrapper headTop = head.top();
     DisiWrapper tailTop = tail.top();
     while (headTop.doc < min) {
@@ -246,8 +220,8 @@ private void scoreWindowMultipleScorers(
 
     if (maxFreq >= minShouldMatch) {
       // There might be matches in other scorers from the tail too
-      for (int i = 0; i < tail.size(); ++i) {
-        leads[maxFreq++] = tail.get(i);
+      for (DisiWrapper disiWrapper : tail) {
+        leads[maxFreq++] = disiWrapper;
       }
       tail.clear();
 

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

@@ -18,6 +18,7 @@
 
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Objects;
 import org.apache.lucene.util.Bits;
@@ -51,12 +52,7 @@ private static class BulkScorerAndNext {
       throw new IllegalArgumentException();
     }
     this.scorers =
-        new PriorityQueue<>(scorers.size()) {
-          @Override
-          protected boolean lessThan(BulkScorerAndNext a, BulkScorerAndNext b) {
-            return a.next < b.next;
-          }
-        };
+        PriorityQueue.usingComparator(scorers.size(), Comparator.comparingInt(b -> b.next));
     for (BulkScorer scorer : scorers) {
       this.scorers.add(new BulkScorerAndNext(scorer));
     }

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

@@ -20,6 +20,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.Iterator;
 import java.util.List;
 import org.apache.lucene.index.Impact;
@@ -260,12 +261,8 @@ public List<Impact> getImpacts(int level) {
             final int docIdUpTo = getDocIdUpTo(level);
 
             PriorityQueue<SubIterator> pq =
-                new PriorityQueue<>(impacts.length) {
-                  @Override
-                  protected boolean lessThan(SubIterator a, SubIterator b) {
-                    return a.current.freq < b.current.freq;
-                  }
-                };
+                PriorityQueue.usingComparator(
+                    impacts.length, Comparator.comparingInt(si -> si.current.freq));
 
             boolean hasImpacts = false;
             List<Impact> onlyImpactList = null;

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

@@ -20,6 +20,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -412,7 +413,7 @@ private boolean termArraysEquals(Term[][] termArrays1, Term[][] termArrays2) {
    */
   public static class UnionPostingsEnum extends PostingsEnum {
     /** queue ordered by docid */
-    final DocsQueue docsQueue;
+    final PriorityQueue<PostingsEnum> docsQueue;
 
     /** cost of this enum: sum of its subs */
     final long cost;
@@ -427,7 +428,8 @@ public static class UnionPostingsEnum extends PostingsEnum {
     final PostingsEnum[] subs;
 
     public UnionPostingsEnum(Collection<PostingsEnum> subs) {
-      docsQueue = new DocsQueue(subs.size());
+      docsQueue =
+          PriorityQueue.usingComparator(subs.size(), Comparator.comparingInt(PostingsEnum::docID));
       long cost = 0;
       for (PostingsEnum sub : subs) {
         docsQueue.add(sub);
@@ -511,18 +513,6 @@ public BytesRef getPayload() throws IOException {
       return null; // payloads are unsupported
     }
 
-    /** disjunction of postings ordered by docid. */
-    static class DocsQueue extends PriorityQueue<PostingsEnum> {
-      DocsQueue(int size) {
-        super(size);
-      }
-
-      @Override
-      public final boolean lessThan(PostingsEnum a, PostingsEnum b) {
-        return a.docID() < b.docID();
-      }
-    }
-
     /**
      * queue of terms for a single document. its a sorted array of all the positions from all the
      * postings
@@ -592,12 +582,7 @@ public static class UnionFullPostingsEnum extends UnionPostingsEnum {
     public UnionFullPostingsEnum(List<PostingsEnum> subs) {
       super(subs);
       this.posQueue =
-          new PriorityQueue<PostingsAndPosition>(subs.size()) {
-            @Override
-            protected boolean lessThan(PostingsAndPosition a, PostingsAndPosition b) {
-              return a.pos < b.pos;
-            }
-          };
+          PriorityQueue.usingComparator(subs.size(), Comparator.comparingInt(p -> p.pos));
       this.subs = new ArrayList<>();
       for (PostingsEnum pe : subs) {
         this.subs.add(new PostingsAndPosition(pe));

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

@@ -18,6 +18,7 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Comparator;
 import java.util.List;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.PostingsEnum;
@@ -58,12 +59,8 @@ protected WeightOrDocIdSetIterator rewriteInner(
           throws IOException {
         DocIdSetBuilder otherTerms = new DocIdSetBuilder(context.reader().maxDoc(), terms);
         PriorityQueue<PostingsEnum> highFrequencyTerms =
-            new PriorityQueue<>(collectedTerms.size()) {
-              @Override
-              protected boolean lessThan(PostingsEnum a, PostingsEnum b) {
-                return a.cost() < b.cost();
-              }
-            };
+            PriorityQueue.usingComparator(
+                collectedTerms.size(), Comparator.comparingLong(PostingsEnum::cost));
 
         // Handle the already-collected terms:
         PostingsEnum reuse = null;

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

@@ -17,6 +17,7 @@
 package org.apache.lucene.search;
 
 import java.io.IOException;
+import java.util.Comparator;
 import java.util.stream.LongStream;
 import java.util.stream.StreamSupport;
 import org.apache.lucene.codecs.lucene103.Lucene103PostingsFormat;
@@ -49,12 +50,7 @@ static long costWithMinShouldMatch(LongStream costs, int numScorers, int minShou
     // If we recurse infinitely, we find out that the cost of a msm query is the sum of the
     // costs of the num_scorers - minShouldMatch + 1 least costly scorers
     final PriorityQueue<Long> pq =
-        new PriorityQueue<Long>(numScorers - minShouldMatch + 1) {
-          @Override
-          protected boolean lessThan(Long a, Long b) {
-            return a > b;
-          }
-        };
+        PriorityQueue.usingComparator(numScorers - minShouldMatch + 1, Comparator.reverseOrder());
     costs.forEach(pq::insertWithOverflow);
     return StreamSupport.stream(pq.spliterator(), false).mapToLong(Number::longValue).sum();
   }