jimczi commented on a change in pull request #913: LUCENE-8995: TopSuggestDocsCollector#collect should be able to signal rejection URL: https://github.com/apache/lucene-solr/pull/913#discussion_r330667847
########## File path: lucene/suggest/src/test/org/apache/lucene/search/suggest/document/TestPrefixCompletionQuery.java ########## @@ -253,6 +263,126 @@ public void testDocFiltering() throws Exception { iw.close(); } + /** + * Test that the correct amount of documents are collected if using a collector that also rejects documents. + */ + public void testCollectorThatRejects() throws Exception { + // use synonym analyzer to have multiple paths to same suggested document. This mock adds "dog" as synonym for "dogs" + Analyzer analyzer = new MockSynonymAnalyzer(); + RandomIndexWriter iw = new RandomIndexWriter(random(), dir, iwcWithSuggestField(analyzer, "suggest_field")); + List<Entry> expectedResults = new ArrayList<Entry>(); + + for (int docCount = 10; docCount > 0; docCount--) { + Document document = new Document(); + String value = "ab" + docCount + " dogs"; + document.add(new SuggestField("suggest_field", value, docCount)); + expectedResults.add(new Entry(value, docCount)); + iw.addDocument(document); + } + + if (rarely()) { + iw.commit(); + } + + DirectoryReader reader = iw.getReader(); + SuggestIndexSearcher indexSearcher = new SuggestIndexSearcher(reader); + + PrefixCompletionQuery query = new PrefixCompletionQuery(analyzer, new Term("suggest_field", "ab")); + int topN = 5; + + // use a TopSuggestDocsCollector that rejects results with duplicate docIds + TopSuggestDocsCollector collector = new TopSuggestDocsCollector(topN, false) { + + private Set<Integer> seenDocIds = new HashSet<>(); + + @Override + public boolean collect(int docID, CharSequence key, CharSequence context, float score) throws IOException { + int globalDocId = docID + docBase; + boolean collected = false; + if (seenDocIds.contains(globalDocId) == false) { + super.collect(docID, key, context, score); + seenDocIds.add(globalDocId); + collected = true; + } + return collected; + } + + @Override + protected boolean canReject() { + return true; + } + }; + + indexSearcher.suggest(query, collector); + TopSuggestDocs suggestions = collector.get(); + assertSuggestions(suggestions, expectedResults.subList(0, topN).toArray(new Entry[0])); + assertTrue(suggestions.isComplete()); + + reader.close(); + iw.close(); + } + + /** + * A large scale tests where the collector rejects based on docIds + */ + public void testCollectorWithManyRejects() throws Exception { + Analyzer analyzer = new MockAnalyzer(random()); + RandomIndexWriter iw = new RandomIndexWriter(random(), dir, iwcWithSuggestField(analyzer, "suggest_field")); + Set<Integer> acceptedDocs = new HashSet<>(); + List<Entry> expectedResults = new ArrayList<Entry>(); + + for (int docCount = 0; docCount < 10000; docCount++) { + Document document = new Document(); + String value = "ab" + RandomStrings.randomAsciiAlphanumOfLength(random(), 10) +"_" + docCount; + document.add(new SuggestField("suggest_field", value, docCount)); + if (random().nextDouble() > 0.75) { Review comment: the maximum queue size is `5000` so we should ensure that we don't reject more than this number if we want to ensure that the search is complete. If you change the live docs to contain at least `5000` docs, this test should work fine. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org With regards, Apache Git Services --------------------------------------------------------------------- To unsubscribe, e-mail: issues-unsubscr...@lucene.apache.org For additional commands, e-mail: issues-h...@lucene.apache.org