public void testCollector() throws IOException {
TotalHitCountCollector collector = new TotalHitCountCollector();
ProfileCollector profileCollector = new ProfileCollector(collector);
assertEquals(0, profileCollector.getTime());
final LeafCollector leafCollector = profileCollector.getLeafCollector(reader.leaves().get(0));
assertThat(profileCollector.getTime(), greaterThan(0L));
long time = profileCollector.getTime();
leafCollector.setScorer(null);
assertThat(profileCollector.getTime(), greaterThan(time));
time = profileCollector.getTime();
leafCollector.collect(0);
assertThat(profileCollector.getTime(), greaterThan(time));
}
private void collectHit(LeafCollector collector, DocsAndCost[] dims) throws IOException {
// if (DEBUG) {
// System.out.println(" hit");
// }
collector.collect(collectDocID);
if (drillDownCollector != null) {
drillDownLeafCollector.collect(collectDocID);
}
// TODO: we could "fix" faceting of the sideways counts
// to do this "union" (of the drill down hits) in the
// end instead:
// Tally sideways counts:
for (DocsAndCost dim : dims) {
dim.sidewaysLeafCollector.collect(collectDocID);
}
}
@Override
protected void search(List<LeafReaderContext> leaves, Weight weight, Collector collector)
throws IOException {
for (LeafReaderContext ctx : leaves) { // search each subreader
// we force the use of Scorer (not BulkScorer) to make sure
// that the scorer passed to LeafCollector.setScorer supports
// Scorer.getChildren
Scorer scorer = weight.scorer(ctx);
if (scorer != null) {
final LeafCollector leafCollector = collector.getLeafCollector(ctx);
leafCollector.setScorer(scorer);
final Bits liveDocs = ctx.reader().getLiveDocs();
for (int doc = scorer.nextDoc();
doc != DocIdSetIterator.NO_MORE_DOCS;
doc = scorer.nextDoc()) {
if (liveDocs == null || liveDocs.get(doc)) {
leafCollector.collect(doc);
}
}
}
}
}
@Test
public void testRandom() throws Exception {
Directory directory = newDirectory();
final Random r = random();
final IndexWriterConfig iwc =
LuceneTestCase.newIndexWriterConfig(r, new MockAnalyzer(r))
.setMaxBufferedDocs(IndexWriterConfig.DISABLE_AUTO_FLUSH)
.setRAMBufferSizeMB(
scaledRandomIntBetween(16, 64)); // we might index a lot - don't go crazy here
RandomIndexWriter indexWriter = new RandomIndexWriter(r, directory, iwc);
int numUniqueChildValues = scaledRandomIntBetween(100, 2000);
String[] childValues = new String[numUniqueChildValues];
for (int i = 0; i < numUniqueChildValues; i++) {
childValues[i] = Integer.toString(i);
}
IntOpenHashSet filteredOrDeletedDocs = new IntOpenHashSet();
int childDocId = 0;
int numParentDocs = scaledRandomIntBetween(1, numUniqueChildValues);
ObjectObjectOpenHashMap<String, NavigableMap<String, FloatArrayList>> childValueToParentIds =
new ObjectObjectOpenHashMap<>();
for (int parentDocId = 0; parentDocId < numParentDocs; parentDocId++) {
boolean markParentAsDeleted = rarely();
boolean filterMe = rarely();
String parent = Integer.toString(parentDocId);
Document document = new Document();
document.add(
new StringField(UidFieldMapper.NAME, Uid.createUid("parent", parent), Field.Store.YES));
document.add(new StringField(TypeFieldMapper.NAME, "parent", Field.Store.NO));
if (markParentAsDeleted) {
filteredOrDeletedDocs.add(parentDocId);
document.add(new StringField("delete", "me", Field.Store.NO));
}
if (filterMe) {
filteredOrDeletedDocs.add(parentDocId);
document.add(new StringField("filter", "me", Field.Store.NO));
}
indexWriter.addDocument(document);
int numChildDocs = scaledRandomIntBetween(0, 100);
for (int i = 0; i < numChildDocs; i++) {
boolean markChildAsDeleted = rarely();
String childValue = childValues[random().nextInt(childValues.length)];
document = new Document();
document.add(
new StringField(
UidFieldMapper.NAME,
Uid.createUid("child", Integer.toString(childDocId++)),
Field.Store.NO));
document.add(new StringField(TypeFieldMapper.NAME, "child", Field.Store.NO));
document.add(
new StringField(
ParentFieldMapper.NAME, Uid.createUid("parent", parent), Field.Store.NO));
document.add(new StringField("field1", childValue, Field.Store.NO));
if (markChildAsDeleted) {
document.add(new StringField("delete", "me", Field.Store.NO));
}
indexWriter.addDocument(document);
if (!markChildAsDeleted) {
NavigableMap<String, FloatArrayList> parentIdToChildScores;
if (childValueToParentIds.containsKey(childValue)) {
parentIdToChildScores = childValueToParentIds.lget();
} else {
childValueToParentIds.put(childValue, parentIdToChildScores = new TreeMap<>());
}
if (!markParentAsDeleted && !filterMe) {
FloatArrayList childScores = parentIdToChildScores.get(parent);
if (childScores == null) {
parentIdToChildScores.put(parent, childScores = new FloatArrayList());
}
childScores.add(1f);
}
}
}
}
// Delete docs that are marked to be deleted.
indexWriter.deleteDocuments(new Term("delete", "me"));
indexWriter.commit();
IndexReader indexReader = DirectoryReader.open(directory);
IndexSearcher searcher = new IndexSearcher(indexReader);
Engine.Searcher engineSearcher =
new Engine.Searcher(ChildrenQueryTests.class.getSimpleName(), searcher);
((TestSearchContext) SearchContext.current())
.setSearcher(new ContextIndexSearcher(SearchContext.current(), engineSearcher));
int max = numUniqueChildValues / 4;
for (int i = 0; i < max; i++) {
// Simulate a parent update
if (random().nextBoolean()) {
final int numberOfUpdatableParents = numParentDocs - filteredOrDeletedDocs.size();
int numberOfUpdates =
RandomInts.randomIntBetween(
random(), 0, Math.min(numberOfUpdatableParents, TEST_NIGHTLY ? 25 : 5));
for (int j = 0; j < numberOfUpdates; j++) {
int parentId;
do {
parentId = random().nextInt(numParentDocs);
} while (filteredOrDeletedDocs.contains(parentId));
String parentUid = Uid.createUid("parent", Integer.toString(parentId));
indexWriter.deleteDocuments(new Term(UidFieldMapper.NAME, parentUid));
Document document = new Document();
document.add(new StringField(UidFieldMapper.NAME, parentUid, Field.Store.YES));
document.add(new StringField(TypeFieldMapper.NAME, "parent", Field.Store.NO));
indexWriter.addDocument(document);
}
indexReader.close();
indexReader = DirectoryReader.open(indexWriter.w, true);
searcher = new IndexSearcher(indexReader);
engineSearcher =
new Engine.Searcher(ChildrenConstantScoreQueryTests.class.getSimpleName(), searcher);
((TestSearchContext) SearchContext.current())
.setSearcher(new ContextIndexSearcher(SearchContext.current(), engineSearcher));
}
String childValue = childValues[random().nextInt(numUniqueChildValues)];
int shortCircuitParentDocSet = random().nextInt(numParentDocs);
ScoreType scoreType = ScoreType.values()[random().nextInt(ScoreType.values().length)];
// leave min/max set to 0 half the time
int minChildren = random().nextInt(2) * scaledRandomIntBetween(0, 110);
int maxChildren = random().nextInt(2) * scaledRandomIntBetween(minChildren, 110);
QueryBuilder queryBuilder =
hasChildQuery("child", constantScoreQuery(termQuery("field1", childValue)))
.scoreType(scoreType.name().toLowerCase(Locale.ENGLISH))
.minChildren(minChildren)
.maxChildren(maxChildren)
.setShortCircuitCutoff(shortCircuitParentDocSet);
// Using a FQ, will invoke / test the Scorer#advance(..) and also let the Weight#scorer not
// get live docs as acceptedDocs
queryBuilder = filteredQuery(queryBuilder, notFilter(termFilter("filter", "me")));
Query query = parseQuery(queryBuilder);
BitSetCollector collector = new BitSetCollector(indexReader.maxDoc());
int numHits = 1 + random().nextInt(25);
TopScoreDocCollector actualTopDocsCollector = TopScoreDocCollector.create(numHits);
searcher.search(query, MultiCollector.wrap(collector, actualTopDocsCollector));
FixedBitSet actualResult = collector.getResult();
FixedBitSet expectedResult = new FixedBitSet(indexReader.maxDoc());
TopScoreDocCollector expectedTopDocsCollector = TopScoreDocCollector.create(numHits);
if (childValueToParentIds.containsKey(childValue)) {
LeafReader slowLeafReader = SlowCompositeReaderWrapper.wrap(indexReader);
final FloatArrayList[] scores = new FloatArrayList[slowLeafReader.maxDoc()];
Terms terms = slowLeafReader.terms(UidFieldMapper.NAME);
if (terms != null) {
NavigableMap<String, FloatArrayList> parentIdToChildScores = childValueToParentIds.lget();
TermsEnum termsEnum = terms.iterator(null);
DocsEnum docsEnum = null;
for (Map.Entry<String, FloatArrayList> entry : parentIdToChildScores.entrySet()) {
int count = entry.getValue().elementsCount;
if (count >= minChildren && (maxChildren == 0 || count <= maxChildren)) {
TermsEnum.SeekStatus seekStatus =
termsEnum.seekCeil(Uid.createUidAsBytes("parent", entry.getKey()));
if (seekStatus == TermsEnum.SeekStatus.FOUND) {
docsEnum =
termsEnum.docs(slowLeafReader.getLiveDocs(), docsEnum, DocsEnum.FLAG_NONE);
expectedResult.set(docsEnum.nextDoc());
scores[docsEnum.docID()] = new FloatArrayList(entry.getValue());
} else if (seekStatus == TermsEnum.SeekStatus.END) {
break;
}
}
}
}
MockScorer mockScorer = new MockScorer(scoreType);
final LeafCollector leafCollector =
expectedTopDocsCollector.getLeafCollector(slowLeafReader.getContext());
leafCollector.setScorer(mockScorer);
for (int doc = expectedResult.nextSetBit(0);
doc < slowLeafReader.maxDoc();
doc =
doc + 1 >= expectedResult.length()
? DocIdSetIterator.NO_MORE_DOCS
: expectedResult.nextSetBit(doc + 1)) {
mockScorer.scores = scores[doc];
leafCollector.collect(doc);
}
}
assertBitSet(actualResult, expectedResult, searcher);
assertTopDocs(actualTopDocsCollector.topDocs(), expectedTopDocsCollector.topDocs());
}
indexWriter.close();
indexReader.close();
directory.close();
}
/**
* Accumulates groups for the BlockJoinQuery specified by its slot.
*
* @param slot Search query's slot
* @param offset Parent docs offset
* @param maxDocsPerGroup Upper bound of documents per group number
* @param withinGroupOffset Offset within each group of child docs
* @param withinGroupSort Sort criteria within groups
* @param fillSortFields Specifies whether to add sort fields or not
* @return TopGroups for the query specified by slot
* @throws IOException if there is a low-level I/O error
*/
@SuppressWarnings({"unchecked", "rawtypes"})
private TopGroups<Integer> accumulateGroups(
int slot,
int offset,
int maxDocsPerGroup,
int withinGroupOffset,
Sort withinGroupSort,
boolean fillSortFields)
throws IOException {
final GroupDocs<Integer>[] groups = new GroupDocs[sortedGroups.length - offset];
final FakeScorer fakeScorer = new FakeScorer();
int totalGroupedHitCount = 0;
// System.out.println("slot=" + slot);
for (int groupIDX = offset; groupIDX < sortedGroups.length; groupIDX++) {
final OneGroup og = sortedGroups[groupIDX];
final int numChildDocs;
if (slot == -1 || slot >= og.counts.length) {
numChildDocs = 0;
} else {
numChildDocs = og.counts[slot];
}
// Number of documents in group should be bounded to prevent redundant memory allocation
final int numDocsInGroup = Math.max(1, Math.min(numChildDocs, maxDocsPerGroup));
// System.out.println("parent doc=" + og.doc + " numChildDocs=" + numChildDocs + " maxDocsPG="
// + maxDocsPerGroup);
// At this point we hold all docs w/ in each group,
// unsorted; we now sort them:
final TopDocsCollector<?> collector;
if (withinGroupSort == null) {
// System.out.println("sort by score");
// Sort by score
if (!trackScores) {
throw new IllegalArgumentException(
"cannot sort by relevance within group: trackScores=false");
}
collector = TopScoreDocCollector.create(numDocsInGroup);
} else {
// Sort by fields
collector =
TopFieldCollector.create(
withinGroupSort, numDocsInGroup, fillSortFields, trackScores, trackMaxScore);
}
LeafCollector leafCollector = collector.getLeafCollector(og.readerContext);
leafCollector.setScorer(fakeScorer);
for (int docIDX = 0; docIDX < numChildDocs; docIDX++) {
// System.out.println("docIDX=" + docIDX + " vs " + og.docs[slot].length);
final int doc = og.docs[slot][docIDX];
fakeScorer.doc = doc;
if (trackScores) {
fakeScorer.score = og.scores[slot][docIDX];
}
leafCollector.collect(doc);
}
totalGroupedHitCount += numChildDocs;
final Object[] groupSortValues;
if (fillSortFields) {
groupSortValues = new Object[comparators.length];
for (int sortFieldIDX = 0; sortFieldIDX < comparators.length; sortFieldIDX++) {
groupSortValues[sortFieldIDX] = comparators[sortFieldIDX].value(og.slot);
}
} else {
groupSortValues = null;
}
final TopDocs topDocs = collector.topDocs(withinGroupOffset, numDocsInGroup);
groups[groupIDX - offset] =
new GroupDocs<>(
og.score,
topDocs.getMaxScore(),
numChildDocs,
topDocs.scoreDocs,
og.doc,
groupSortValues);
}
return new TopGroups<>(
new TopGroups<>(
sort.getSort(),
withinGroupSort == null ? null : withinGroupSort.getSort(),
0,
totalGroupedHitCount,
groups,
maxScore),
totalHitCount);
}
@Override
public int score(LeafCollector collector, Bits acceptDocs, int min, int maxDoc)
throws IOException {
if (min != 0) {
throw new IllegalArgumentException("min must be 0, got " + min);
}
if (maxDoc != Integer.MAX_VALUE) {
throw new IllegalArgumentException("maxDoc must be Integer.MAX_VALUE");
}
// if (DEBUG) {
// System.out.println("\nscore: reader=" + context.reader());
// }
// System.out.println("score r=" + context.reader());
FakeScorer scorer = new FakeScorer();
collector.setScorer(scorer);
if (drillDownCollector != null) {
drillDownLeafCollector = drillDownCollector.getLeafCollector(context);
drillDownLeafCollector.setScorer(scorer);
} else {
drillDownLeafCollector = null;
}
for (DocsAndCost dim : dims) {
dim.sidewaysLeafCollector = dim.sidewaysCollector.getLeafCollector(context);
dim.sidewaysLeafCollector.setScorer(scorer);
}
// TODO: if we ever allow null baseScorer ... it will
// mean we DO score docs out of order ... hmm, or if we
// change up the order of the conjuntions below
assert baseScorer != null;
// some scorers, eg ReqExlScorer, can hit NPE if cost is called after nextDoc
long baseQueryCost = baseScorer.cost();
final int numDims = dims.length;
long drillDownCost = 0;
for (int dim = 0; dim < numDims; dim++) {
drillDownCost += dims[dim].approximation.cost();
}
long drillDownAdvancedCost = 0;
if (numDims > 1) {
drillDownAdvancedCost = dims[1].approximation.cost();
}
// Position all scorers to their first matching doc:
baseScorer.nextDoc();
for (DocsAndCost dim : dims) {
dim.approximation.nextDoc();
}
/*
System.out.println("\nbaseDocID=" + baseScorer.docID() + " est=" + estBaseHitCount);
System.out.println(" maxDoc=" + context.reader().maxDoc());
System.out.println(" maxCost=" + maxCost);
System.out.println(" dims[0].freq=" + dims[0].freq);
if (numDims > 1) {
System.out.println(" dims[1].freq=" + dims[1].freq);
}
*/
if (scoreSubDocsAtOnce || baseQueryCost < drillDownCost / 10) {
// System.out.println("queryFirst: baseScorer=" + baseScorer + " disis.length=" + disis.length
// + " bits.length=" + bits.length);
doQueryFirstScoring(acceptDocs, collector, dims);
} else if (numDims > 1 && drillDownAdvancedCost < baseQueryCost / 10) {
// System.out.println("drillDownAdvance");
doDrillDownAdvanceScoring(acceptDocs, collector, dims);
} else {
// System.out.println("union");
doUnionScoring(acceptDocs, collector, dims);
}
return Integer.MAX_VALUE;
}
private void collectNearMiss(LeafCollector sidewaysCollector) throws IOException {
// if (DEBUG) {
// System.out.println(" missingDim=" + dim);
// }
sidewaysCollector.collect(collectDocID);
}