LUCENE-9565 Fix competitive iteration (#1952)

PR #1351 introduced a sort optimization where documents can be skipped.
But iteration over competitive iterators was not properly organized,
as they were not storing the current docID, and
when competitive iterator was updated the current doc ID was lost.

This patch fixed it.

Relates to #1351

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

@@ -204,9 +204,17 @@ public int score(LeafCollector collector, Bits acceptDocs, int min, int max) thr
       collector.setScorer(scorer);
       DocIdSetIterator scorerIterator = twoPhase == null ? iterator : twoPhase.approximation();
       DocIdSetIterator collectorIterator = collector.competitiveIterator();
-      // if possible filter scorerIterator to keep only competitive docs as defined by collector
-      DocIdSetIterator filteredIterator = collectorIterator == null ? scorerIterator :
-          ConjunctionDISI.intersectIterators(Arrays.asList(scorerIterator, collectorIterator));
+      DocIdSetIterator filteredIterator;
+      if (collectorIterator == null) {
+        filteredIterator = scorerIterator;
+      } else {
+        if (scorerIterator.docID() != -1) {
+          // Wrap ScorerIterator to start from -1 for conjunction 
+          scorerIterator = new RangeDISIWrapper(scorerIterator, max);
+        }
+        // filter scorerIterator to keep only competitive docs as defined by collector
+        filteredIterator = ConjunctionDISI.intersectIterators(Arrays.asList(scorerIterator, collectorIterator));
+      }
       if (filteredIterator.docID() == -1 && min == 0 && max == DocIdSetIterator.NO_MORE_DOCS) {
         scoreAll(collector, filteredIterator, twoPhase, acceptDocs);
         return DocIdSetIterator.NO_MORE_DOCS;
@@ -266,4 +274,45 @@ static void scoreAll(LeafCollector collector, DocIdSetIterator iterator, TwoPhas
     }
   }
 
+  /**
+   * Wraps an internal docIdSetIterator for it to start with docID = -1
+   */
+  protected static class RangeDISIWrapper extends DocIdSetIterator {
+    private final DocIdSetIterator in;
+    private final int min;
+    private final int max;
+    private int docID = -1;
+
+    public RangeDISIWrapper(DocIdSetIterator in, int max) {
+      this.in = in;
+      this.min = in.docID();
+      this.max = max;
+    }
+
+    @Override
+    public int docID() {
+      return docID;
+    }
+
+    @Override
+    public int nextDoc() throws IOException {
+      return advance(docID + 1);
+    }
+
+    @Override
+    public int advance(int target) throws IOException {
+      target = Math.max(min, target);
+      if (target >= max) {
+        return docID = NO_MORE_DOCS;
+      }
+      return docID = in.advance(target);
+    }
+
+    @Override
+    public long cost() {
+      return Math.min(max - min, in.cost());
+    }
+
+  }
+
 }

lucene/core/src/java/org/apache/lucene/search/comparators/DocComparator.java

@@ -133,14 +133,16 @@ public DocIdSetIterator competitiveIterator() {
                 return null;
             } else {
                 return new DocIdSetIterator() {
+                    private int docID = -1;
+
                     @Override
                     public int nextDoc() throws IOException {
-                        return competitiveIterator.nextDoc();
+                        return advance(docID + 1);
                     }
 
                     @Override
                     public int docID() {
-                        return competitiveIterator.docID();
+                        return docID;
                     }
 
                     @Override
@@ -150,7 +152,7 @@ public long cost() {
 
                     @Override
                     public int advance(int target) throws IOException {
-                        return competitiveIterator.advance(target);
+                        return docID = competitiveIterator.advance(target);
                     }
                 };
             }

lucene/core/src/java/org/apache/lucene/search/comparators/NumericComparator.java

@@ -220,14 +220,16 @@ public PointValues.Relation compare(byte[] minPackedValue, byte[] maxPackedValue
         public DocIdSetIterator competitiveIterator() {
             if (enableSkipping == false) return null;
             return new DocIdSetIterator() {
+                private int docID = -1;
+
                 @Override
                 public int nextDoc() throws IOException {
-                    return competitiveIterator.nextDoc();
+                    return advance(docID + 1);
                 }
 
                 @Override
                 public int docID() {
-                    return competitiveIterator.docID();
+                    return docID;
                 }
 
                 @Override
@@ -237,7 +239,7 @@ public long cost() {
 
                 @Override
                 public int advance(int target) throws IOException {
-                    return competitiveIterator.advance(target);
+                    return docID = competitiveIterator.advance(target);
                 }
             };
         }