GITHUB#14399: decide index-sort early termination per segment in TopFieldCollector

lucene/CHANGES.txt

@@ -7,6 +7,9 @@ http://s.apache.org/luceneversions
 
 API Changes
 ---------------------
+* GITHUB#14399: Add a default LeafFieldComparator#disableSkipping method so document skipping can be
+  disabled per segment. (Serhiy Bzhezytskyy)
+
 * GITHUB#15929: Rename CollectionStatistics to FieldStats and TermStatistics to TermStats. (Zhou Hui)
 
 * GITHUB#15763: Deprecate Operations.complement() method. This operation can be slow and is not
@@ -190,6 +193,10 @@ Optimizations
 
 Bug Fixes
 ---------------------
+* GITHUB#14399: TopFieldCollector now decides per segment whether the search sort is a prefix of the
+  index sort, instead of caching the decision from the first segment. This fixes incorrect results
+  when searching a MultiReader that combines indexes with different index sorts. (Serhiy Bzhezytskyy)
+
 * GITHUB#14049: Randomize KNN codec params in RandomCodec. Fixes scalar quantization div-by-zero
   when all values are identical. (Mike Sokolov)
 

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

@@ -111,4 +111,13 @@ default DocIdSetIterator competitiveIterator() throws IOException {
    * collector when hits threshold is reached.
    */
   default void setHitsThresholdReached() throws IOException {}
+
+  /**
+   * Informs this leaf comparator that document skipping should be disabled for this segment. Called
+   * by {@link TopFieldCollector} when the search sort is a prefix of <em>this segment's</em> index
+   * sort, so the collector can already terminate early and any extra skipping work in the
+   * comparator is redundant. This is a per-segment decision: a {@link
+   * org.apache.lucene.index.MultiReader} may combine segments with different index sorts.
+   */
+  default void disableSkipping() {}
 }

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

@@ -45,20 +45,18 @@ private abstract class TopFieldLeafCollector implements LeafCollector {
 
     final LeafFieldComparator comparator;
     final int reverseMul;
+    // Whether the search sort is a prefix of this segment's index sort (decided per segment).
+    final boolean searchSortPartOfIndexSort;
     Scorable scorer;
     boolean collectedAllCompetitiveHits = false;
 
     TopFieldLeafCollector(FieldValueHitQueue<Entry> queue, Sort sort, LeafReaderContext context)
         throws IOException {
-      // as all segments are sorted in the same way, enough to check only the 1st segment for
-      // indexSort
-      if (searchSortPartOfIndexSort == null) {
-        final Sort indexSort = context.reader().getMetaData().sort();
-        searchSortPartOfIndexSort = canEarlyTerminate(sort, indexSort);
-        if (searchSortPartOfIndexSort) {
-          firstComparator.disableSkipping();
-        }
-      }
+      // Whether the search sort is a prefix of the index sort is decided per segment: a MultiReader
+      // may combine segments with different index sorts, so this cannot be cached across leaves
+      // (GITHUB#14399).
+      final Sort indexSort = context.reader().getMetaData().sort();
+      searchSortPartOfIndexSort = canEarlyTerminate(sort, indexSort);
       LeafFieldComparator[] comparators = queue.getComparators(context);
       int[] reverseMuls = queue.getReverseMul();
       if (comparators.length == 1) {
@@ -68,6 +66,12 @@ private abstract class TopFieldLeafCollector implements LeafCollector {
         this.reverseMul = 1;
         this.comparator = new MultiLeafFieldComparator(comparators, reverseMuls);
       }
+      if (searchSortPartOfIndexSort) {
+        // Early termination handles this segment; skipping work in the comparator is redundant.
+        for (LeafFieldComparator comparator : comparators) {
+          comparator.disableSkipping();
+        }
+      }
     }
 
     void countHit() throws IOException {
@@ -304,8 +308,6 @@ public void collect(int doc) throws IOException {
   final FieldComparator<?> firstComparator;
   final boolean canSetMinScore;
 
-  Boolean searchSortPartOfIndexSort = null; // shows if Search Sort if a part of the Index Sort
-
   // an accumulator that maintains the maximum of the segment's minimum competitive scores
   final MaxScoreAccumulator minScoreAcc;
   // the current local minimum competitive score already propagated to the underlying scorer

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

@@ -98,14 +98,22 @@ public void disableSkipping() {
   public abstract class NumericLeafComparator implements LeafFieldComparator {
     private final LeafReaderContext context;
     protected final NumericDocValues docValues;
-    private final CompetitiveDISIBuilder competitiveDISIBuilder;
+    private CompetitiveDISIBuilder competitiveDISIBuilder;
 
     public NumericLeafComparator(LeafReaderContext context) throws IOException {
       this.context = context;
       this.docValues = getNumericDocValues(context, field);
       this.competitiveDISIBuilder = buildCompetitiveDISIBuilder();
     }
 
+    @Override
+    public void disableSkipping() {
+      // Drop the competitive iterator so this segment is scanned without skipping; the collector
+      // has determined the search sort is a prefix of this segment's index sort and will terminate
+      // early on its own.
+      this.competitiveDISIBuilder = null;
+    }
+
     protected CompetitiveDISIBuilder buildCompetitiveDISIBuilder() throws IOException {
       if (pruning == Pruning.NONE) {
         return null;

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

@@ -214,7 +214,7 @@ class TermOrdValLeafComparator implements LeafFieldComparator {
     /** Which ordinal to use for a missing value. */
     final int missingOrd;
 
-    private final CompetitiveState competitiveState;
+    private CompetitiveState competitiveState;
 
     private final boolean dense;
 
@@ -491,6 +491,13 @@ private void updateCompetitiveIterator() throws IOException {
       competitiveState.update(minOrd, maxOrd);
     }
 
+    @Override
+    public void disableSkipping() {
+      // Drop the competitive iterator so this segment is scanned without skipping; the collector
+      // will terminate early on its own because the search sort is a prefix of this segment's sort.
+      competitiveState = null;
+    }
+
     @Override
     public DocIdSetIterator competitiveIterator() {
       return competitiveState != null ? competitiveState.iterator : null;

lucene/core/src/test/org/apache/lucene/search/TestTopFieldCollectorEarlyTermination.java

@@ -27,8 +27,11 @@
 import org.apache.lucene.document.Field.Store;
 import org.apache.lucene.document.NumericDocValuesField;
 import org.apache.lucene.document.StringField;
+import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.MultiReader;
 import org.apache.lucene.index.SerialMergeScheduler;
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.IndexSearcher.LeafReaderContextPartition;
@@ -41,6 +44,7 @@
 import org.apache.lucene.tests.util.LuceneTestCase;
 import org.apache.lucene.tests.util.TestUtil;
 import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.IOUtils;
 
 public class TestTopFieldCollectorEarlyTermination extends LuceneTestCase {
 
@@ -264,4 +268,86 @@ public void testCanEarlyTerminateOnPrefix() {
                 new SortField("c", SortField.Type.LONG),
                 new SortField("b", SortField.Type.STRING))));
   }
+
+  /**
+   * GITHUB#14399: TopFieldCollector caches whether the search sort is a prefix of the index sort
+   * after inspecting only the first leaf. That is safe for a single index (IndexWriter enforces one
+   * index sort), but a MultiReader can span indexes with different index sorts. Here the first
+   * index is sorted so the search sort IS a prefix (early termination is eligible) while the second
+   * index is sorted the opposite way (it is NOT). The cached "yes" wrongly early-terminates the
+   * second leaf and drops results that should rank first.
+   */
+  public void testMultiReaderWithDifferentIndexSorts() throws IOException {
+    final Sort ascSort = new Sort(new SortField("ndv", SortField.Type.LONG));
+
+    // Index A: sorted ndv ASC, many docs with a moderate value (50). Under an ASC search sort this
+    // leaf is prefix-sorted, so the collector caches "search sort is part of index sort" = true and
+    // calls disableSkipping().
+    Directory dirA = newDirectory();
+    IndexWriterConfig iwcA = newIndexWriterConfig().setIndexSort(ascSort);
+    try (RandomIndexWriter w = new RandomIndexWriter(random(), dirA, iwcA)) {
+      for (int i = 0; i < 20; i++) {
+        Document doc = new Document();
+        doc.add(new NumericDocValuesField("ndv", 50L));
+        w.addDocument(doc);
+      }
+      w.forceMerge(1);
+    }
+
+    // Index B: sorted ndv DESC (a DIFFERENT index sort). Its most competitive doc for an ASC search
+    // (value 1) is written LAST, so in docid order the leaf is [90, 90, ..., 1]. If leaf B is
+    // wrongly
+    // treated as prefix-sorted (leaf A's cached decision), collection terminates in docid order
+    // after
+    // the threshold and never reaches the trailing value 1 -- the true top result is dropped.
+    Directory dirB = newDirectory();
+    IndexWriterConfig iwcB =
+        newIndexWriterConfig()
+            .setIndexSort(new Sort(new SortField("ndv", SortField.Type.LONG, true)));
+    try (RandomIndexWriter w = new RandomIndexWriter(random(), dirB, iwcB)) {
+      for (int i = 0; i < 20; i++) {
+        Document doc = new Document();
+        doc.add(new NumericDocValuesField("ndv", 90L));
+        w.addDocument(doc);
+      }
+      Document winner = new Document();
+      winner.add(new NumericDocValuesField("ndv", 1L)); // the smallest value overall
+      w.addDocument(winner);
+      w.forceMerge(1);
+    }
+
+    DirectoryReader readerA = DirectoryReader.open(dirA);
+    DirectoryReader readerB = DirectoryReader.open(dirB);
+    // MultiReader with A first, so the first leaf is the ASC-sorted one (search sort is a prefix ->
+    // eligible for early termination); the B leaf is DESC-sorted (not a prefix).
+    MultiReader multiReader = new MultiReader(readerA, readerB);
+    try {
+      // Force both leaves into a single slice so one collector (with one shared
+      // searchSortPartOfIndexSort cache) sees both the ASC-sorted and DESC-sorted leaves.
+      IndexSearcher searcher =
+          new IndexSearcher(multiReader) {
+            @Override
+            protected LeafSlice[] slices(List<LeafReaderContext> leaves) {
+              List<LeafReaderContextPartition> partitions = new ArrayList<>();
+              for (LeafReaderContext ctx : leaves) {
+                partitions.add(LeafReaderContextPartition.createForEntireSegment(ctx));
+              }
+              return new LeafSlice[] {new LeafSlice(partitions)};
+            }
+          };
+      // numHits=1, threshold=1: want the single smallest value overall, which is 1 (last docid of
+      // leaf B). Correct answer is 1; the bug drops it and returns 50 (from leaf A).
+      TopFieldCollectorManager manager = new TopFieldCollectorManager(ascSort, 1, null, 1);
+      TopFieldDocs td = searcher.search(new MatchAllDocsQuery(), manager);
+
+      assertEquals(1, td.scoreDocs.length);
+      long topValue = (long) ((FieldDoc) td.scoreDocs[0]).fields[0];
+      assertEquals("the smallest value (1, trailing docid of leaf B) must win", 1L, topValue);
+    } finally {
+      multiReader.close();
+      readerA.close();
+      readerB.close();
+      IOUtils.close(dirA, dirB);
+    }
+  }
 }