review feedback + add ut

This commit is contained in:
jimczi 2019-09-30 13:14:48 +02:00
parent 58fabbed2b
commit 75ff81c18f
3 changed files with 259 additions and 6 deletions

View File

@ -143,7 +143,7 @@ public abstract class TopFieldCollector extends TopDocsCollector<Entry> {
} 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<Entry> {
} 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<Entry> {
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<Entry> {
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<Entry> {
}
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<Entry> {
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<Entry> {
* and a shared bottom value checker to propagate the minimum score accross segments if
* the primary sort is by relevancy.
*/
public static CollectorManager<TopFieldCollector, TopFieldDocs> createSharedManager(Sort sort, int numHits, FieldDoc after, int totalHitsThreshold) {
public static CollectorManager<TopFieldCollector, TopFieldDocs> createSharedManager(Sort sort, int numHits, FieldDoc after,
int totalHitsThreshold) {
return new CollectorManager<>() {
private final HitsThresholdChecker hitsThresholdChecker = HitsThresholdChecker.createShared(totalHitsThreshold);

View File

@ -76,7 +76,8 @@ public abstract class TopScoreDocCollector extends TopDocsCollector<ScoreDoc> {
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<ScoreDoc> {
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> {
}
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<ScoreDoc> {
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<ScoreDoc> {
}
scorer.setMinCompetitiveScore(bottomScore);
minScore = bottomScore;
totalHitsRelation = TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO;
}
}

View File

@ -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<TopScoreDocCollector, TopDocs> 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<TopFieldCollector, TopFieldDocs> 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();
}
}