diff --git a/lucene/core/src/java/org/apache/lucene/search/TopFieldCollector.java b/lucene/core/src/java/org/apache/lucene/search/TopFieldCollector.java index d49ecf8eac4..2f62ba04ac5 100644 --- a/lucene/core/src/java/org/apache/lucene/search/TopFieldCollector.java +++ b/lucene/core/src/java/org/apache/lucene/search/TopFieldCollector.java @@ -143,7 +143,7 @@ public abstract class TopFieldCollector extends TopDocsCollector { } else { collectedAllCompetitiveHits = true; } - } else if (totalHitsRelation == Relation.EQUAL_TO) { + } else if (totalHitsRelation == Relation.EQUAL_TO || shouldUpdateMinScore()) { // we just reached totalHitsThreshold, we can start setting the min // competitive score now updateMinCompetitiveScore(scorer); @@ -238,7 +238,7 @@ public abstract class TopFieldCollector extends TopDocsCollector { } else { collectedAllCompetitiveHits = true; } - } else if (totalHitsRelation == Relation.EQUAL_TO) { + } else if (totalHitsRelation == Relation.EQUAL_TO || shouldUpdateMinScore()) { // we just reached totalHitsThreshold, we can start setting the min // competitive score now updateMinCompetitiveScore(scorer); @@ -290,6 +290,10 @@ public abstract class TopFieldCollector extends TopDocsCollector { final BottomValueChecker bottomValueChecker; final FieldComparator.RelevanceComparator firstComparator; final boolean canSetMinScore; + // the minimum score (if canSetMinScore is true) that is currently used by the scorer, + // can be different than pqTop.score if the provided bottomValueChecker reports a minimum + // score that is greater than the local one. + float minScore; final int numComparators; FieldValueHitQueue.Entry bottom = null; boolean queueFull; @@ -331,6 +335,10 @@ public abstract class TopFieldCollector extends TopDocsCollector { return scoreMode; } + protected boolean shouldUpdateMinScore() { + return bottomValueChecker != null ? bottomValueChecker.getBottomValue() > minScore : false; + } + protected void updateMinCompetitiveScore(Scorable scorer) throws IOException { if (canSetMinScore && hitsThresholdChecker.isThresholdReached() && (queueFull || (bottomValueChecker != null && bottomValueChecker.getBottomValue() > 0f))) { @@ -347,6 +355,7 @@ public abstract class TopFieldCollector extends TopDocsCollector { } assert maxMinScore > 0f; scorer.setMinCompetitiveScore(maxMinScore); + minScore = maxMinScore; totalHitsRelation = TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO; } } @@ -406,7 +415,7 @@ public abstract class TopFieldCollector extends TopDocsCollector { throw new IllegalArgumentException("totalHitsThreshold must be >= 0, got " + totalHitsThreshold); } - return create(sort, numHits, after, HitsThresholdChecker.create(totalHitsThreshold), null); + return create(sort, numHits, after, HitsThresholdChecker.create(totalHitsThreshold), null /* bottomValueChecker */); } /** @@ -449,7 +458,8 @@ public abstract class TopFieldCollector extends TopDocsCollector { * and a shared bottom value checker to propagate the minimum score accross segments if * the primary sort is by relevancy. */ - public static CollectorManager createSharedManager(Sort sort, int numHits, FieldDoc after, int totalHitsThreshold) { + public static CollectorManager createSharedManager(Sort sort, int numHits, FieldDoc after, + int totalHitsThreshold) { return new CollectorManager<>() { private final HitsThresholdChecker hitsThresholdChecker = HitsThresholdChecker.createShared(totalHitsThreshold); diff --git a/lucene/core/src/java/org/apache/lucene/search/TopScoreDocCollector.java b/lucene/core/src/java/org/apache/lucene/search/TopScoreDocCollector.java index 10c4fc4c43a..a08e969579f 100644 --- a/lucene/core/src/java/org/apache/lucene/search/TopScoreDocCollector.java +++ b/lucene/core/src/java/org/apache/lucene/search/TopScoreDocCollector.java @@ -76,7 +76,8 @@ public abstract class TopScoreDocCollector extends TopDocsCollector { hitsThresholdChecker.incrementHitCount(); if (score <= pqTop.score) { - if (totalHitsRelation == TotalHits.Relation.EQUAL_TO && hitsThresholdChecker.isThresholdReached()) { + if ((totalHitsRelation == TotalHits.Relation.EQUAL_TO || shouldUpdateMinScore()) + && hitsThresholdChecker.isThresholdReached()) { // we just reached totalHitsThreshold, we can start setting the min // competitive score now updateMinCompetitiveScore(scorer); @@ -139,7 +140,8 @@ public abstract class TopScoreDocCollector extends TopDocsCollector { if (score > after.score || (score == after.score && doc <= afterDoc)) { // hit was collected on a previous page - if (totalHitsRelation == TotalHits.Relation.EQUAL_TO && hitsThresholdChecker.isThresholdReached()) { + if ((totalHitsRelation == TotalHits.Relation.EQUAL_TO || shouldUpdateMinScore()) + && hitsThresholdChecker.isThresholdReached()) { // we just reached totalHitsThreshold, we can start setting the min // competitive score now updateMinCompetitiveScore(scorer); @@ -247,6 +249,10 @@ public abstract class TopScoreDocCollector extends TopDocsCollector { } ScoreDoc pqTop; + // the minimum score that is currently used by the scorer, can be different than pqTop.score + // if the provided bottomValueChecker reports a minimum score that is greater than the local + // one. + float minScore; final HitsThresholdChecker hitsThresholdChecker; final BottomValueChecker bottomValueChecker; @@ -277,6 +283,10 @@ public abstract class TopScoreDocCollector extends TopDocsCollector { return hitsThresholdChecker.scoreMode(); } + protected boolean shouldUpdateMinScore() { + return bottomValueChecker != null ? bottomValueChecker.getBottomValue() > minScore : false; + } + protected void updateMinCompetitiveScore(Scorable scorer) throws IOException { if (hitsThresholdChecker.isThresholdReached() && ((bottomValueChecker != null && bottomValueChecker.getBottomValue() > 0) @@ -301,6 +311,7 @@ public abstract class TopScoreDocCollector extends TopDocsCollector { } scorer.setMinCompetitiveScore(bottomScore); + minScore = bottomScore; totalHitsRelation = TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO; } } diff --git a/lucene/core/src/test/org/apache/lucene/search/TestTopDocsCollector.java b/lucene/core/src/test/org/apache/lucene/search/TestTopDocsCollector.java index e0818e2f120..1476b189a70 100644 --- a/lucene/core/src/test/org/apache/lucene/search/TestTopDocsCollector.java +++ b/lucene/core/src/test/org/apache/lucene/search/TestTopDocsCollector.java @@ -476,4 +476,236 @@ public class TestTopDocsCollector extends LuceneTestCase { dir.close(); } + public void testConcurrentMinScoreTopScoreDocs() throws Exception { + Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, newIndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE)); + Document doc = new Document(); + w.addDocuments(Arrays.asList(doc, doc, doc, doc, doc)); + w.flush(); + w.addDocuments(Arrays.asList(doc, doc, doc, doc, doc, doc)); + w.flush(); + w.addDocuments(Arrays.asList(doc, doc)); + w.flush(); + IndexReader reader = DirectoryReader.open(w); + assertEquals(3, reader.leaves().size()); + w.close(); + + CollectorManager manager = + TopScoreDocCollector.createSharedManager(2, null, 0); + TopScoreDocCollector collector = manager.newCollector(); + TopScoreDocCollector collector2 = manager.newCollector(); + assertTrue(collector.bottomValueChecker == collector2.bottomValueChecker); + BottomValueChecker minValueChecker = collector.bottomValueChecker; + + ScoreAndDoc scorer = new ScoreAndDoc(); + ScoreAndDoc scorer2 = new ScoreAndDoc(); + + LeafCollector leafCollector = collector.getLeafCollector(reader.leaves().get(0)); + leafCollector.setScorer(scorer); + LeafCollector leafCollector2 = collector2.getLeafCollector(reader.leaves().get(1)); + leafCollector2.setScorer(scorer2); + + scorer.doc = 0; + scorer.score = 3; + leafCollector.collect(0); + + scorer2.doc = 0; + scorer2.score = 6; + leafCollector2.collect(0); + + scorer.doc = 1; + scorer.score = 2; + leafCollector.collect(1); + assertEquals(minValueChecker.getBottomValue(), 2f, 0f); + assertEquals(scorer.minCompetitiveScore, Math.nextUp(2f), 0f); + assertNull(scorer2.minCompetitiveScore); + + scorer2.doc = 1; + scorer2.score = 9; + leafCollector2.collect(1); + assertEquals(minValueChecker.getBottomValue(), 6f, 0f); + assertEquals(scorer.minCompetitiveScore, Math.nextUp(2f), 0f); + assertEquals(scorer2.minCompetitiveScore, Math.nextUp(6f), 0f); + + scorer2.doc = 2; + scorer2.score = 7; + leafCollector2.collect(2); + assertEquals(minValueChecker.getBottomValue(), 7f, 0f); + assertEquals(scorer.minCompetitiveScore, Math.nextUp(2f), 0f); + assertEquals(scorer2.minCompetitiveScore, Math.nextUp(7f), 0f); + + scorer2.doc = 3; + scorer2.score = 1; + leafCollector2.collect(3); + assertEquals(minValueChecker.getBottomValue(), 7f, 0f); + assertEquals(scorer.minCompetitiveScore, Math.nextUp(2f), 0f); + assertEquals(scorer2.minCompetitiveScore, Math.nextUp(7f), 0f); + + scorer.doc = 2; + scorer.score = 10; + leafCollector.collect(2); + assertEquals(minValueChecker.getBottomValue(), 7f, 0f); + assertEquals(scorer.minCompetitiveScore, 7f, 0f); + assertEquals(scorer2.minCompetitiveScore, Math.nextUp(7f), 0f); + + scorer.doc = 3; + scorer.score = 11; + leafCollector.collect(3); + assertEquals(minValueChecker.getBottomValue(), 10, 0f); + assertEquals(scorer.minCompetitiveScore, Math.nextUp(10f), 0f); + assertEquals(scorer2.minCompetitiveScore, Math.nextUp(7f), 0f); + + TopScoreDocCollector collector3 = manager.newCollector(); + LeafCollector leafCollector3 = collector3.getLeafCollector(reader.leaves().get(2)); + ScoreAndDoc scorer3 = new ScoreAndDoc(); + leafCollector3.setScorer(scorer3); + assertEquals(scorer3.minCompetitiveScore, 10f, 0f); + + scorer3.doc = 0; + scorer3.score = 1f; + leafCollector3.collect(0); + assertEquals(minValueChecker.getBottomValue(), 10f, 0f); + assertEquals(scorer3.minCompetitiveScore, 10f, 0f); + + scorer.doc = 4; + scorer.score = 11; + leafCollector.collect(4); + assertEquals(minValueChecker.getBottomValue(), 11f, 0f); + assertEquals(scorer.minCompetitiveScore, Math.nextUp(11f), 0f); + assertEquals(scorer2.minCompetitiveScore, Math.nextUp(7f), 0f); + assertEquals(scorer3.minCompetitiveScore, 10f, 0f); + + scorer3.doc = 1; + scorer3.score = 2f; + leafCollector3.collect(1); + assertEquals(minValueChecker.getBottomValue(), 11f, 0f); + assertEquals(scorer.minCompetitiveScore, Math.nextUp(11f), 0f); + assertEquals(scorer2.minCompetitiveScore, Math.nextUp(7f), 0f); + assertEquals(scorer3.minCompetitiveScore, 11f, 0f); + + + TopDocs topDocs = manager.reduce(Arrays.asList(collector, collector2, collector3)); + assertEquals(11, topDocs.totalHits.value); + assertEquals(new TotalHits(11, TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO), topDocs.totalHits); + + reader.close(); + dir.close(); + } + + public void testConcurrentMinScoreTopFieldDocs() throws Exception { + Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, newIndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE)); + Document doc = new Document(); + w.addDocuments(Arrays.asList(doc, doc, doc, doc, doc)); + w.flush(); + w.addDocuments(Arrays.asList(doc, doc, doc, doc, doc, doc)); + w.flush(); + w.addDocuments(Arrays.asList(doc, doc)); + w.flush(); + IndexReader reader = DirectoryReader.open(w); + assertEquals(3, reader.leaves().size()); + w.close(); + + Sort sort = new Sort(SortField.FIELD_SCORE, SortField.FIELD_DOC); + CollectorManager manager = + TopFieldCollector.createSharedManager(sort, 2, null, 0); + TopFieldCollector collector = manager.newCollector(); + TopFieldCollector collector2 = manager.newCollector(); + assertTrue(collector.bottomValueChecker == collector2.bottomValueChecker); + BottomValueChecker minValueChecker = collector.bottomValueChecker; + + ScoreAndDoc scorer = new ScoreAndDoc(); + ScoreAndDoc scorer2 = new ScoreAndDoc(); + + LeafCollector leafCollector = collector.getLeafCollector(reader.leaves().get(0)); + leafCollector.setScorer(scorer); + LeafCollector leafCollector2 = collector2.getLeafCollector(reader.leaves().get(1)); + leafCollector2.setScorer(scorer2); + + scorer.doc = 0; + scorer.score = 3; + leafCollector.collect(0); + + scorer2.doc = 0; + scorer2.score = 6; + leafCollector2.collect(0); + + scorer.doc = 1; + scorer.score = 2; + leafCollector.collect(1); + assertEquals(minValueChecker.getBottomValue(), 2f, 0f); + assertEquals(scorer.minCompetitiveScore, 2f, 0f); + assertNull(scorer2.minCompetitiveScore); + + scorer2.doc = 1; + scorer2.score = 9; + leafCollector2.collect(1); + assertEquals(minValueChecker.getBottomValue(), 6f, 0f); + assertEquals(scorer.minCompetitiveScore, 2f, 0f); + assertEquals(scorer2.minCompetitiveScore, 6f, 0f); + + scorer2.doc = 2; + scorer2.score = 7; + leafCollector2.collect(2); + assertEquals(minValueChecker.getBottomValue(), 7f, 0f); + assertEquals(scorer.minCompetitiveScore, 2f, 0f); + assertEquals(scorer2.minCompetitiveScore, 7f, 0f); + + scorer2.doc = 3; + scorer2.score = 1; + leafCollector2.collect(3); + assertEquals(minValueChecker.getBottomValue(), 7f, 0f); + assertEquals(scorer.minCompetitiveScore, 2f, 0f); + assertEquals(scorer2.minCompetitiveScore, 7f, 0f); + + scorer.doc = 2; + scorer.score = 10; + leafCollector.collect(2); + assertEquals(minValueChecker.getBottomValue(), 7f, 0f); + assertEquals(scorer.minCompetitiveScore, 7f, 0f); + assertEquals(scorer2.minCompetitiveScore, 7f, 0f); + + scorer.doc = 3; + scorer.score = 11; + leafCollector.collect(3); + assertEquals(minValueChecker.getBottomValue(), 10, 0f); + assertEquals(scorer.minCompetitiveScore, 10f, 0f); + assertEquals(scorer2.minCompetitiveScore, 7f, 0f); + + TopFieldCollector collector3 = manager.newCollector(); + LeafCollector leafCollector3 = collector3.getLeafCollector(reader.leaves().get(2)); + ScoreAndDoc scorer3 = new ScoreAndDoc(); + leafCollector3.setScorer(scorer3); + assertEquals(scorer3.minCompetitiveScore, 10f, 0f); + + scorer3.doc = 0; + scorer3.score = 1f; + leafCollector3.collect(0); + assertEquals(minValueChecker.getBottomValue(), 10f, 0f); + assertEquals(scorer3.minCompetitiveScore, 10f, 0f); + + scorer.doc = 4; + scorer.score = 11; + leafCollector.collect(4); + assertEquals(minValueChecker.getBottomValue(), 11f, 0f); + assertEquals(scorer.minCompetitiveScore, 11f, 0f); + assertEquals(scorer2.minCompetitiveScore, 7f, 0f); + assertEquals(scorer3.minCompetitiveScore, 10f, 0f); + + scorer3.doc = 1; + scorer3.score = 2f; + leafCollector3.collect(1); + assertEquals(minValueChecker.getBottomValue(), 11f, 0f); + assertEquals(scorer.minCompetitiveScore, 11f, 0f); + assertEquals(scorer2.minCompetitiveScore, 7f, 0f); + assertEquals(scorer3.minCompetitiveScore, 11f, 0f); + + + TopFieldDocs topDocs = manager.reduce(Arrays.asList(collector, collector2, collector3)); + assertEquals(11, topDocs.totalHits.value); + assertEquals(new TotalHits(11, TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO), topDocs.totalHits); + + reader.close(); + dir.close(); + } }