LUCENE-10002: replace more usages of search(Query, Collector) in tests (#787)

This commit replaces more usages of search(Query, Collector) with calling the corresponding search(Query, CollectorManager) instead. This round focuses on tests that implement custom collector, that need a corresponding collector manager.
This commit is contained in:
Luca Cavanna 2022-04-06 11:06:10 +02:00 committed by GitHub
parent 74e9716aec
commit 1cf1b301af
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
8 changed files with 163 additions and 143 deletions

View File

@ -20,7 +20,6 @@ import com.carrotsearch.randomizedtesting.generators.RandomPicks;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.BitSet;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
@ -42,9 +41,11 @@ import org.apache.lucene.search.similarities.ClassicSimilarity;
import org.apache.lucene.store.Directory;
import org.apache.lucene.tests.analysis.MockAnalyzer;
import org.apache.lucene.tests.index.RandomIndexWriter;
import org.apache.lucene.tests.search.FixedBitSetCollector;
import org.apache.lucene.tests.search.QueryUtils;
import org.apache.lucene.tests.util.LuceneTestCase;
import org.apache.lucene.tests.util.TestUtil;
import org.apache.lucene.util.FixedBitSet;
import org.apache.lucene.util.IOUtils;
import org.apache.lucene.util.NamedThreadFactory;
import org.apache.lucene.util.automaton.Operations;
@ -416,30 +417,8 @@ public class TestBooleanQuery extends LuceneTestCase {
dir.close();
}
private static BitSet getMatches(IndexSearcher searcher, Query query) throws IOException {
BitSet set = new BitSet();
searcher.search(
query,
new SimpleCollector() {
int docBase = 0;
@Override
public ScoreMode scoreMode() {
return ScoreMode.COMPLETE_NO_SCORES;
}
@Override
protected void doSetNextReader(LeafReaderContext context) throws IOException {
super.doSetNextReader(context);
docBase = context.docBase;
}
@Override
public void collect(int doc) throws IOException {
set.set(docBase + doc);
}
});
return set;
private static FixedBitSet getMatches(IndexSearcher searcher, Query query) throws IOException {
return searcher.search(query, FixedBitSetCollector.createManager(searcher.reader.maxDoc()));
}
public void testFILTERClauseBehavesLikeMUST() throws IOException {
@ -473,8 +452,8 @@ public class TestBooleanQuery extends LuceneTestCase {
bq2.add(q, Occur.FILTER);
}
final BitSet matches1 = getMatches(searcher, bq1.build());
final BitSet matches2 = getMatches(searcher, bq2.build());
final FixedBitSet matches1 = getMatches(searcher, bq1.build());
final FixedBitSet matches2 = getMatches(searcher, bq2.build());
assertEquals(matches1, matches2);
}

View File

@ -19,6 +19,7 @@ package org.apache.lucene.search;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
@ -99,7 +100,7 @@ public class TestMultiCollector extends LuceneTestCase {
}
final IndexReader reader = w.getReader();
w.close();
final IndexSearcher searcher = newSearcher(reader);
final IndexSearcher searcher = newSearcher(reader, true, true, false);
Map<TotalHitCountCollector, Integer> expectedCounts = new HashMap<>();
List<Collector> collectors = new ArrayList<>();
final int numCollectors = TestUtil.nextInt(random(), 1, 5);
@ -110,7 +111,19 @@ public class TestMultiCollector extends LuceneTestCase {
expectedCounts.put(collector, expectedCount);
collectors.add(new TerminateAfterCollector(collector, terminateAfter));
}
searcher.search(new MatchAllDocsQuery(), MultiCollector.wrap(collectors));
searcher.search(
new MatchAllDocsQuery(),
new CollectorManager<Collector, Void>() {
@Override
public Collector newCollector() {
return MultiCollector.wrap(collectors);
}
@Override
public Void reduce(Collection<Collector> collectors) {
return null;
}
});
for (Map.Entry<TotalHitCountCollector, Integer> expectedCount : expectedCounts.entrySet()) {
assertEquals(expectedCount.getValue().intValue(), expectedCount.getKey().getTotalHits());
}

View File

@ -18,6 +18,7 @@ package org.apache.lucene.search;
import java.io.IOException;
import java.util.BitSet;
import java.util.Collection;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.index.DirectoryReader;
@ -77,6 +78,34 @@ public class TestScorerPerf extends LuceneTestCase {
return sets;
}
private static final class CountingHitCollectorManager
implements CollectorManager<CountingHitCollector, CountingHitCollector> {
private final boolean validate;
private final FixedBitSet result;
CountingHitCollectorManager(boolean validate, FixedBitSet result) {
this.validate = validate;
this.result = result;
}
@Override
public CountingHitCollector newCollector() {
return validate ? new MatchingHitCollector(result) : new CountingHitCollector();
}
@Override
public CountingHitCollector reduce(Collection<CountingHitCollector> collectors)
throws IOException {
CountingHitCollector result = new CountingHitCollector();
for (CountingHitCollector collector : collectors) {
result.count += collector.count;
result.sum += collector.sum;
}
return result;
}
}
public static class CountingHitCollector extends SimpleCollector {
int count = 0;
int sum = 0;
@ -191,10 +220,8 @@ public class TestScorerPerf extends LuceneTestCase {
for (int j = 0; j < nClauses; j++) {
result = addClause(sets, bq, result);
}
CountingHitCollector hc =
validate ? new MatchingHitCollector(result) : new CountingHitCollector();
s.search(bq.build(), hc);
s.search(bq.build(), new CountingHitCollectorManager(validate, result));
ret += hc.getSum();
if (validate) assertEquals(result.cardinality(), hc.getCount());
@ -227,8 +254,7 @@ public class TestScorerPerf extends LuceneTestCase {
} // outer
CountingHitCollector hc =
validate ? new MatchingHitCollector(result) : new CountingHitCollector();
s.search(oq.build(), hc);
s.search(oq.build(), new CountingHitCollectorManager(validate, result));
nMatches += hc.getCount();
ret += hc.getSum();
if (validate) assertEquals(result.cardinality(), hc.getCount());
@ -259,8 +285,7 @@ public class TestScorerPerf extends LuceneTestCase {
bq.add(tq, BooleanClause.Occur.MUST);
}
CountingHitCollector hc = new CountingHitCollector();
s.search(bq.build(), hc);
CountingHitCollector hc = s.search(bq.build(), new CountingHitCollectorManager(false, null));
nMatches += hc.getCount();
ret += hc.getSum();
}
@ -301,8 +326,7 @@ public class TestScorerPerf extends LuceneTestCase {
oq.add(bq.build(), BooleanClause.Occur.MUST);
} // outer
CountingHitCollector hc = new CountingHitCollector();
s.search(oq.build(), hc);
CountingHitCollector hc = s.search(oq.build(), new CountingHitCollectorManager(false, null));
nMatches += hc.getCount();
ret += hc.getSum();
}
@ -325,8 +349,7 @@ public class TestScorerPerf extends LuceneTestCase {
builder.setSlop(termsInIndex);
PhraseQuery q = builder.build();
CountingHitCollector hc = new CountingHitCollector();
s.search(q, hc);
CountingHitCollector hc = s.search(q, new CountingHitCollectorManager(false, null));
ret += hc.getSum();
}

View File

@ -18,6 +18,7 @@ package org.apache.lucene.search;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
@ -78,6 +79,34 @@ public class TestSubScorerFreqs extends LuceneTestCase {
dir = null;
}
private static class CountingCollectorManager
implements CollectorManager<CountingCollector, Map<Integer, Map<Query, Float>>> {
private final Set<String> relationships;
CountingCollectorManager(Set<String> relationships) {
this.relationships = relationships;
}
@Override
public CountingCollector newCollector() {
TopScoreDocCollector topScoreDocCollector =
TopScoreDocCollector.create(10, Integer.MAX_VALUE);
return relationships == null
? new CountingCollector(topScoreDocCollector)
: new CountingCollector(topScoreDocCollector, relationships);
}
@Override
public Map<Integer, Map<Query, Float>> reduce(Collection<CountingCollector> collectors) {
Map<Integer, Map<Query, Float>> docCounts = new HashMap<>();
for (CountingCollector collector : collectors) {
docCounts.putAll(collector.docCounts);
}
return docCounts;
}
}
private static class CountingCollector extends FilterCollector {
public final Map<Integer, Map<Query, Float>> docCounts = new HashMap<>();
@ -135,16 +164,15 @@ public class TestSubScorerFreqs extends LuceneTestCase {
@Test
public void testTermQuery() throws Exception {
TermQuery q = new TermQuery(new Term("f", "d"));
CountingCollector c = new CountingCollector(TopScoreDocCollector.create(10, Integer.MAX_VALUE));
s.search(q, c);
Map<Integer, Map<Query, Float>> docCounts = s.search(q, new CountingCollectorManager(null));
final int maxDocs = s.getIndexReader().maxDoc();
assertEquals(maxDocs, c.docCounts.size());
assertEquals(maxDocs, docCounts.size());
for (int i = 0; i < maxDocs; i++) {
Map<Query, Float> doc0 = c.docCounts.get(i);
Map<Query, Float> doc0 = docCounts.get(i);
assertEquals(1, doc0.size());
assertEquals(4.0F, doc0.get(q), FLOAT_TOLERANCE);
Map<Query, Float> doc1 = c.docCounts.get(++i);
Map<Query, Float> doc1 = docCounts.get(++i);
assertEquals(1, doc1.size());
assertEquals(1.0F, doc1.get(q), FLOAT_TOLERANCE);
}
@ -174,14 +202,13 @@ public class TestSubScorerFreqs extends LuceneTestCase {
Collections.singleton("MUST"), new HashSet<>(Arrays.asList("MUST", "SHOULD")));
for (final Set<String> occur : occurList) {
CountingCollector c =
new CountingCollector(TopScoreDocCollector.create(10, Integer.MAX_VALUE), occur);
s.search(query.build(), c);
Map<Integer, Map<Query, Float>> docCounts =
s.search(query.build(), new CountingCollectorManager(occur));
final int maxDocs = s.getIndexReader().maxDoc();
assertEquals(maxDocs, c.docCounts.size());
assertEquals(maxDocs, docCounts.size());
boolean includeOptional = occur.contains("SHOULD");
for (int i = 0; i < maxDocs; i++) {
Map<Query, Float> doc0 = c.docCounts.get(i);
Map<Query, Float> doc0 = docCounts.get(i);
// Y doesnt exist in the index, so it's not in the scorer tree
assertEquals(4, doc0.size());
assertEquals(1.0F, doc0.get(aQuery), FLOAT_TOLERANCE);
@ -190,7 +217,7 @@ public class TestSubScorerFreqs extends LuceneTestCase {
assertEquals(3.0F, doc0.get(cQuery), FLOAT_TOLERANCE);
}
Map<Query, Float> doc1 = c.docCounts.get(++i);
Map<Query, Float> doc1 = docCounts.get(++i);
// Y doesnt exist in the index, so it's not in the scorer tree
assertEquals(4, doc1.size());
assertEquals(1.0F, doc1.get(aQuery), FLOAT_TOLERANCE);
@ -205,16 +232,15 @@ public class TestSubScorerFreqs extends LuceneTestCase {
@Test
public void testPhraseQuery() throws Exception {
PhraseQuery q = new PhraseQuery("f", "b", "c");
CountingCollector c = new CountingCollector(TopScoreDocCollector.create(10, Integer.MAX_VALUE));
s.search(q, c);
Map<Integer, Map<Query, Float>> docCounts = s.search(q, new CountingCollectorManager(null));
final int maxDocs = s.getIndexReader().maxDoc();
assertEquals(maxDocs, c.docCounts.size());
assertEquals(maxDocs, docCounts.size());
for (int i = 0; i < maxDocs; i++) {
Map<Query, Float> doc0 = c.docCounts.get(i);
Map<Query, Float> doc0 = docCounts.get(i);
assertEquals(1, doc0.size());
assertEquals(2.0F, doc0.get(q), FLOAT_TOLERANCE);
Map<Query, Float> doc1 = c.docCounts.get(++i);
Map<Query, Float> doc1 = docCounts.get(++i);
assertEquals(1, doc1.size());
assertEquals(1.0F, doc1.get(q), FLOAT_TOLERANCE);
}

View File

@ -20,6 +20,7 @@ import static org.apache.lucene.search.SortField.FIELD_SCORE;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collection;
import java.util.Comparator;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
@ -423,44 +424,57 @@ public class TestTopFieldCollector extends LuceneTestCase {
.add(bar, Occur.SHOULD)
.add(baz, Occur.SHOULD)
.build();
final IndexSearcher searcher = new IndexSearcher(reader);
final IndexSearcher searcher = newSearcher(reader);
for (Sort sort :
new Sort[] {new Sort(FIELD_SCORE), new Sort(new SortField("f", SortField.Type.SCORE))}) {
final TopFieldCollector topCollector =
TopFieldCollector.create(sort, TestUtil.nextInt(random(), 1, 2), Integer.MAX_VALUE);
final Collector assertingCollector =
new Collector() {
int numHits = TestUtil.nextInt(random(), 1, 2);
searcher.search(
query,
new CollectorManager<Collector, Void>() {
@Override
public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
final LeafCollector in = topCollector.getLeafCollector(context);
return new FilterLeafCollector(in) {
public Collector newCollector() {
return new Collector() {
final TopFieldCollector topCollector =
TopFieldCollector.create(sort, numHits, Integer.MAX_VALUE);
@Override
public void setScorer(final Scorable scorer) throws IOException {
Scorable s =
new FilterScorable(scorer) {
public LeafCollector getLeafCollector(LeafReaderContext context)
throws IOException {
final LeafCollector in = topCollector.getLeafCollector(context);
return new FilterLeafCollector(in) {
@Override
public void setScorer(final Scorable scorer) throws IOException {
Scorable s =
new FilterScorable(scorer) {
int lastComputedDoc = -1;
int lastComputedDoc = -1;
@Override
public float score() throws IOException {
if (lastComputedDoc == docID()) {
throw new AssertionError("Score computed twice on " + docID());
}
lastComputedDoc = docID();
return scorer.score();
}
};
super.setScorer(s);
@Override
public float score() throws IOException {
if (lastComputedDoc == docID()) {
throw new AssertionError("Score computed twice on " + docID());
}
lastComputedDoc = docID();
return scorer.score();
}
};
super.setScorer(s);
}
};
}
@Override
public ScoreMode scoreMode() {
return ScoreMode.COMPLETE;
}
};
}
@Override
public ScoreMode scoreMode() {
return topCollector.scoreMode();
public Void reduce(Collection<Collector> collectors) {
return null;
}
};
searcher.search(query, assertingCollector);
});
}
reader.close();
w.close();

View File

@ -38,10 +38,10 @@ public class TestTotalHitCountCollector extends LuceneTestCase {
IndexReader reader = writer.getReader();
writer.close();
IndexSearcher searcher = newSearcher(reader);
TotalHitCountCollector c = new TotalHitCountCollector();
searcher.search(new MatchAllDocsQuery(), c);
assertEquals(5, c.getTotalHits());
IndexSearcher searcher = newSearcher(reader, true, true, random().nextBoolean());
TotalHitCountCollectorManager collectorManager = new TotalHitCountCollectorManager();
int totalHits = searcher.search(new MatchAllDocsQuery(), collectorManager);
assertEquals(5, totalHits);
reader.close();
indexStore.close();
}

View File

@ -1,49 +0,0 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.lucene.search;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.StringField;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.store.Directory;
import org.apache.lucene.tests.index.RandomIndexWriter;
import org.apache.lucene.tests.util.LuceneTestCase;
public class TestTotalHitCountCollectorManager extends LuceneTestCase {
public void testBasics() throws Exception {
Directory indexStore = newDirectory();
RandomIndexWriter writer = new RandomIndexWriter(random(), indexStore);
for (int i = 0; i < 5; i++) {
Document doc = new Document();
doc.add(new StringField("string", "a" + i, Field.Store.NO));
doc.add(new StringField("string", "b" + i, Field.Store.NO));
writer.addDocument(doc);
}
IndexReader reader = writer.getReader();
writer.close();
IndexSearcher searcher = newSearcher(reader, true, true, true);
TotalHitCountCollectorManager collectorManager = new TotalHitCountCollectorManager();
int totalHits = searcher.search(new MatchAllDocsQuery(), collectorManager);
assertEquals(5, totalHits);
reader.close();
indexStore.close();
}
}

View File

@ -17,15 +17,19 @@
package org.apache.lucene.queryparser.surround.query;
import java.io.IOException;
import java.util.Collection;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.queryparser.surround.parser.QueryParser;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.CollectorManager;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Scorable;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.search.SimpleCollector;
import org.apache.lucene.util.SetOnce;
import org.junit.Assert;
public class BooleanQueryTestFacade {
@ -121,14 +125,24 @@ public class BooleanQueryTestFacade {
Query query = lq.makeLuceneQueryField(fieldName, qf);
/* if (verbose) System.out.println("Lucene: " + query.toString()); */
TestCollector tc = new TestCollector();
IndexReader reader = DirectoryReader.open(dBase.getDb());
IndexSearcher searcher = new IndexSearcher(reader);
try {
searcher.search(query, tc);
} finally {
reader.close();
SetOnce<TestCollector> collector = new SetOnce<>();
try (IndexReader reader = DirectoryReader.open(dBase.getDb())) {
IndexSearcher searcher = new IndexSearcher(reader);
searcher.search(
query,
new CollectorManager<Collector, Void>() {
@Override
public Collector newCollector() {
collector.set(new TestCollector());
return collector.get();
}
@Override
public Void reduce(Collection<Collector> collectors) {
return null;
}
});
}
tc.checkNrHits();
collector.get().checkNrHits();
}
}