Convert simple two-value queues to using Comparators

Author: thecoop

lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene60/bkd/BKDWriter60.java

@@ -20,6 +20,7 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Comparator;
 import java.util.List;
 import java.util.function.IntFunction;
 import org.apache.lucene.codecs.CodecUtil;
@@ -453,29 +454,14 @@ public Relation compare(byte[] minPackedValue, byte[] maxPackedValue) {
     }
   }
 
-  private static class BKDMergeQueue extends PriorityQueue<MergeReader> {
-    private final int bytesPerDim;
-
-    public BKDMergeQueue(int bytesPerDim, int maxSize) {
-      super(maxSize);
-      this.bytesPerDim = bytesPerDim;
-    }
-
-    @Override
-    public boolean lessThan(MergeReader a, MergeReader b) {
-      assert a != b;
-
-      int cmp =
-          Arrays.compareUnsigned(a.packedValue, 0, bytesPerDim, b.packedValue, 0, bytesPerDim);
-      if (cmp < 0) {
-        return true;
-      } else if (cmp > 0) {
-        return false;
-      }
-
-      // Tie break by sorting smaller docIDs earlier:
-      return a.docID < b.docID;
-    }
+  private static Comparator<MergeReader> mergeComparator(int bytesPerDim) {
+    return ((Comparator<MergeReader>)
+            (a, b) -> {
+              assert a != b;
+              return Arrays.compareUnsigned(
+                  a.packedValue, 0, bytesPerDim, b.packedValue, 0, bytesPerDim);
+            })
+        .thenComparingInt(mr -> mr.docID);
   }
 
   /**
@@ -642,7 +628,8 @@ public long merge(IndexOutput out, List<MergeState.DocMap> docMaps, List<PointVa
       throws IOException {
     assert docMaps == null || readers.size() == docMaps.size();
 
-    BKDMergeQueue queue = new BKDMergeQueue(config.bytesPerDim(), readers.size());
+    PriorityQueue<MergeReader> queue =
+        PriorityQueue.usingComparator(readers.size(), mergeComparator(config.bytesPerDim()));
 
     for (int i = 0; i < readers.size(); i++) {
       PointValues pointValues = readers.get(i);

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

@@ -18,6 +18,7 @@
 
 import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
 
+import java.util.Comparator;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.util.Accountable;
 import org.apache.lucene.util.BytesRef;
@@ -159,24 +160,12 @@ public static Iterator mergedIterator(Iterator[] subs) {
       return subs[0];
     }
 
+    // sort by smaller docID, then larger delGen
     PriorityQueue<Iterator> queue =
-        new PriorityQueue<Iterator>(subs.length) {
-          @Override
-          protected boolean lessThan(Iterator a, Iterator b) {
-            // sort by smaller docID
-            int cmp = Integer.compare(a.docID(), b.docID());
-            if (cmp == 0) {
-              // then by larger delGen
-              cmp = Long.compare(b.delGen(), a.delGen());
-
-              // delGens are unique across our subs:
-              assert cmp != 0;
-            }
-
-            return cmp < 0;
-          }
-        };
-
+        PriorityQueue.usingComparator(
+            subs.length,
+            Comparator.comparingInt(Iterator::docID)
+                .thenComparing(Comparator.comparingLong(Iterator::delGen).reversed()));
     for (Iterator sub : subs) {
       if (sub.nextDoc() != NO_MORE_DOCS) {
         queue.add(sub);

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

@@ -32,6 +32,7 @@
 import org.apache.lucene.internal.hppc.IntHashSet;
 import org.apache.lucene.search.similarities.Similarity.SimScorer;
 import org.apache.lucene.util.FixedBitSet;
+import org.apache.lucene.util.PriorityQueue;
 
 /**
  * Find all slop-valid position-combinations (matches) encountered while traversing/hopping the
@@ -56,7 +57,7 @@ public final class SloppyPhraseMatcher extends PhraseMatcher {
 
   private final int slop;
   private final int numPostings;
-  private final PhraseQueue pq; // for advancing min position
+  private final PriorityQueue<PhrasePositions> pq; // for advancing min position
   private final boolean captureLeadMatch;
 
   private final DocIdSetIterator approximation;
@@ -92,7 +93,14 @@ public SloppyPhraseMatcher(
     this.slop = slop;
     this.numPostings = postings.length;
     this.captureLeadMatch = captureLeadMatch;
-    pq = new PhraseQueue(postings.length);
+    pq =
+        PriorityQueue.usingComparator(
+            postings.length,
+            Comparator.<PhrasePositions>comparingInt(pp -> pp.position)
+                // same doc and pp.position, so decide by actual term positions.
+                // rely on: pp.position == tp.position - offset.
+                .thenComparingInt(pp -> pp.offset)
+                .thenComparingInt(pp -> pp.ord));
     phrasePositions = new PhrasePositions[postings.length];
     for (int i = 0; i < postings.length; ++i) {
       phrasePositions[i] =

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

@@ -16,6 +16,7 @@
  */
 package org.apache.lucene.util;
 
+import java.util.Comparator;
 import java.util.Iterator;
 import java.util.NoSuchElementException;
 
@@ -44,7 +45,7 @@
  */
 public final class MergedIterator<T extends Comparable<T>> implements Iterator<T> {
   private T current;
-  private final TermMergeQueue<T> queue;
+  private final PriorityQueue<SubIterator<T>> queue;
   private final SubIterator<T>[] top;
   private final boolean removeDuplicates;
   private int numTop;
@@ -57,7 +58,11 @@ public MergedIterator(Iterator<T>... iterators) {
   @SuppressWarnings({"unchecked", "rawtypes"})
   public MergedIterator(boolean removeDuplicates, Iterator<T>... iterators) {
     this.removeDuplicates = removeDuplicates;
-    queue = new TermMergeQueue<>(iterators.length);
+    queue =
+        PriorityQueue.usingComparator(
+            iterators.length,
+            Comparator.<SubIterator<T>, T>comparing(it -> it.current)
+                .thenComparingInt(it -> it.index));
     top = new SubIterator[iterators.length];
     int index = 0;
     for (Iterator<T> iterator : iterators) {
@@ -138,21 +143,4 @@ private static class SubIterator<I extends Comparable<I>> {
     I current;
     int index;
   }
-
-  private static class TermMergeQueue<C extends Comparable<C>>
-      extends PriorityQueue<SubIterator<C>> {
-    TermMergeQueue(int size) {
-      super(size);
-    }
-
-    @Override
-    protected boolean lessThan(SubIterator<C> a, SubIterator<C> b) {
-      final int cmp = a.current.compareTo(b.current);
-      if (cmp != 0) {
-        return cmp < 0;
-      } else {
-        return a.index < b.index;
-      }
-    }
-  }
 }

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

@@ -20,6 +20,7 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Comparator;
 import java.util.List;
 import java.util.function.IntFunction;
 import org.apache.lucene.codecs.CodecUtil;
@@ -396,29 +397,14 @@ public Relation compare(byte[] minPackedValue, byte[] maxPackedValue) {
     }
   }
 
-  private static class BKDMergeQueue extends PriorityQueue<MergeReader> {
-    private final ArrayUtil.ByteArrayComparator comparator;
-
-    public BKDMergeQueue(int bytesPerDim, int maxSize) {
-      super(maxSize);
-      this.comparator = ArrayUtil.getUnsignedComparator(bytesPerDim);
-    }
-
-    @Override
-    public boolean lessThan(MergeReader a, MergeReader b) {
-      assert a != b;
-
-      int cmp = comparator.compare(a.packedValue, 0, b.packedValue, 0);
-
-      if (cmp < 0) {
-        return true;
-      } else if (cmp > 0) {
-        return false;
-      }
-
-      // Tie break by sorting smaller docIDs earlier:
-      return a.docID < b.docID;
-    }
+  private static Comparator<MergeReader> mergeComparator(int bytesPerDim) {
+    ByteArrayComparator comparator = ArrayUtil.getUnsignedComparator(bytesPerDim);
+    return ((Comparator<MergeReader>)
+            (a, b) -> {
+              assert a != b;
+              return comparator.compare(a.packedValue, 0, b.packedValue, 0);
+            })
+        .thenComparingInt(mr -> mr.docID); // Tie break by sorting smaller docIDs earlier
   }
 
   /** flat representation of a kd-tree */
@@ -651,7 +637,8 @@ public IORunnable merge(
       throws IOException {
     assert docMaps == null || readers.size() == docMaps.size();
 
-    BKDMergeQueue queue = new BKDMergeQueue(config.bytesPerDim(), readers.size());
+    PriorityQueue<MergeReader> queue =
+        PriorityQueue.usingComparator(readers.size(), mergeComparator(config.bytesPerDim()));
 
     for (int i = 0; i < readers.size(); i++) {
       PointValues pointValues = readers.get(i);

lucene/core/src/java/org/apache/lucene/util/bkd/BKDWriter.java

@@ -21,6 +21,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.List;
 import org.apache.lucene.facet.FacetsCollector.MatchingDocs;
 import org.apache.lucene.index.DocValues;
@@ -419,14 +420,12 @@ public FacetResult getTopChildren(int topN, String dim, String... path) {
       return new FacetResult(field, new String[0], totCount, new LabelAndValue[0], 0);
     }
 
+    // sort by count descending, breaking ties by value ascending:
     PriorityQueue<Entry> pq =
-        new PriorityQueue<>(Math.min(topN, counts.length + hashCounts.size())) {
-          @Override
-          protected boolean lessThan(Entry a, Entry b) {
-            // sort by count descending, breaking ties by value ascending:
-            return a.count < b.count || (a.count == b.count && a.value > b.value);
-          }
-        };
+        PriorityQueue.usingComparator(
+            Math.min(topN, counts.length + hashCounts.size()),
+            Comparator.<Entry>comparingInt(e -> e.count)
+                .thenComparing(Comparator.<Entry>comparingLong(dv -> dv.value).reversed()));
 
     int childCount = 0;
     Entry e = null;

lucene/facet/src/java/org/apache/lucene/facet/LongValueFacetCounts.java

@@ -18,6 +18,7 @@
 
 import java.io.IOException;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.List;
 import org.apache.lucene.facet.FacetCountsWithFilterQuery;
 import org.apache.lucene.facet.FacetResult;
@@ -223,16 +224,10 @@ public FacetResult getTopChildren(int topN, String dim, String... path) throws I
     }
 
     PriorityQueue<Entry> pq =
-        new PriorityQueue<>(Math.min(topN, counts.length)) {
-          @Override
-          protected boolean lessThan(Entry a, Entry b) {
-            int cmp = Integer.compare(a.count, b.count);
-            if (cmp == 0) {
-              cmp = b.label.compareTo(a.label);
-            }
-            return cmp < 0;
-          }
-        };
+        PriorityQueue.usingComparator(
+            Math.min(topN, counts.length),
+            Comparator.<Entry>comparingInt(e -> e.count)
+                .thenComparing(e -> e.label, Comparator.reverseOrder()));
 
     int childCount = 0;
     Entry e = null;

lucene/facet/src/java/org/apache/lucene/facet/range/RangeFacetCounts.java

@@ -18,6 +18,7 @@
 
 import java.io.IOException;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.List;
 import org.apache.lucene.document.BinaryRangeDocValues;
 import org.apache.lucene.document.RangeFieldQuery;
@@ -144,16 +145,10 @@ public FacetResult getTopChildren(int topN, String dim, String... path) throws I
     validateDimAndPathForGetChildren(dim, path);
 
     PriorityQueue<Entry> pq =
-        new PriorityQueue<>(Math.min(topN, counts.length)) {
-          @Override
-          protected boolean lessThan(Entry a, Entry b) {
-            int cmp = Integer.compare(a.count, b.count);
-            if (cmp == 0) {
-              cmp = b.label.compareTo(a.label);
-            }
-            return cmp < 0;
-          }
-        };
+        PriorityQueue.usingComparator(
+            Math.min(topN, counts.length),
+            Comparator.<Entry>comparingInt(e -> e.count)
+                .thenComparing(e -> e.label, Comparator.reverseOrder()));
 
     int childCount = 0;
     Entry e = null;

lucene/facet/src/java/org/apache/lucene/facet/rangeonrange/RangeOnRangeFacetCounts.java

@@ -20,6 +20,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -150,18 +151,10 @@ public List<FacetResult> getTopDims(int topNDims, int topNChildren) throws IOExc
     // Creates priority queue to store top dimensions and sort by their aggregated values/hits and
     // string values.
     PriorityQueue<DimValue> pq =
-        new PriorityQueue<>(topNDims) {
-          @Override
-          protected boolean lessThan(DimValue a, DimValue b) {
-            if (a.value > b.value) {
-              return false;
-            } else if (a.value < b.value) {
-              return true;
-            } else {
-              return a.dim.compareTo(b.dim) > 0;
-            }
-          }
-        };
+        PriorityQueue.usingComparator(
+            topNDims,
+            Comparator.<DimValue>comparingInt(dv -> dv.value)
+                .thenComparing(dv -> dv.dim, Comparator.reverseOrder()));
 
     // Keep track of intermediate results, if we compute them, so we can reuse them later:
     Map<String, TopChildrenForPath> intermediateResults = null;