/**
* Creates a ToParentBlockJoinCollector. The provided sort must not be null. If you pass true
* trackScores, all ToParentBlockQuery instances must not use ScoreMode.None.
*/
public ToParentBlockJoinCollector(
Sort sort, int numParentHits, boolean trackScores, boolean trackMaxScore) throws IOException {
// TODO: allow null sort to be specialized to relevance
// only collector
this.sort = sort;
this.trackMaxScore = trackMaxScore;
if (trackMaxScore) {
maxScore = Float.MIN_VALUE;
}
// System.out.println("numParentHits=" + numParentHits);
this.trackScores = trackScores;
this.numParentHits = numParentHits;
queue = FieldValueHitQueue.create(sort.getSort(), numParentHits);
comparators = queue.getComparators();
}
/**
* Returns the TopGroups for the specified BlockJoinQuery. The groupValue of each GroupDocs will
* be the parent docID for that group. The number of documents within each group is calculated as
* minimum of <code>maxDocsPerGroup</code> and number of matched child documents for that group.
* Returns null if no groups matched.
*
* @param query Search query
* @param withinGroupSort Sort criteria within groups
* @param offset Parent docs offset
* @param maxDocsPerGroup Upper bound of documents per group number
* @param withinGroupOffset Offset within each group of child docs
* @param fillSortFields Specifies whether to add sort fields or not
* @return TopGroups for specified query
* @throws IOException if there is a low-level I/O error
*/
public TopGroups<Integer> getTopGroups(
ToParentBlockJoinQuery query,
Sort withinGroupSort,
int offset,
int maxDocsPerGroup,
int withinGroupOffset,
boolean fillSortFields)
throws IOException {
final Integer _slot = joinQueryID.get(query);
if (_slot == null && totalHitCount == 0) {
return null;
}
if (sortedGroups == null) {
if (offset >= queue.size()) {
return null;
}
sortQueue();
} else if (offset > sortedGroups.length) {
return null;
}
return accumulateGroups(
_slot == null ? -1 : _slot.intValue(),
offset,
maxDocsPerGroup,
withinGroupOffset,
withinGroupSort,
fillSortFields);
}
@Override
public void setNextReader(AtomicReaderContext context) throws IOException {
currentReaderContext = context;
docBase = context.docBase;
for (int compIDX = 0; compIDX < comparators.length; compIDX++) {
queue.setComparator(compIDX, comparators[compIDX].setNextReader(context));
}
}
private void sortQueue() {
sortedGroups = new OneGroup[queue.size()];
for (int downTo = queue.size() - 1; downTo >= 0; downTo--) {
sortedGroups[downTo] = queue.pop();
}
}
@Override
public void collect(int parentDoc) throws IOException {
// System.out.println("\nC parentDoc=" + parentDoc);
totalHitCount++;
float score = Float.NaN;
if (trackMaxScore) {
score = scorer.score();
maxScore = Math.max(maxScore, score);
}
// TODO: we could sweep all joinScorers here and
// aggregate total child hit count, so we can fill this
// in getTopGroups (we wire it to 0 now)
if (queueFull) {
// System.out.println(" queueFull");
// Fastmatch: return if this hit is not competitive
for (int i = 0; ; i++) {
final int c = reverseMul[i] * comparators[i].compareBottom(parentDoc);
if (c < 0) {
// Definitely not competitive.
// System.out.println(" skip");
return;
} else if (c > 0) {
// Definitely competitive.
break;
} else if (i == compEnd) {
// Here c=0. If we're at the last comparator, this doc is not
// competitive, since docs are visited in doc Id order, which means
// this doc cannot compete with any other document in the queue.
// System.out.println(" skip");
return;
}
}
// System.out.println(" competes! doc=" + (docBase + parentDoc));
// This hit is competitive - replace bottom element in queue & adjustTop
for (int i = 0; i < comparators.length; i++) {
comparators[i].copy(bottom.slot, parentDoc);
}
if (!trackMaxScore && trackScores) {
score = scorer.score();
}
bottom.doc = docBase + parentDoc;
bottom.readerContext = currentReaderContext;
bottom.score = score;
copyGroups(bottom);
bottom = queue.updateTop();
for (int i = 0; i < comparators.length; i++) {
comparators[i].setBottom(bottom.slot);
}
} else {
// Startup transient: queue is not yet full:
final int comparatorSlot = totalHitCount - 1;
// Copy hit into queue
for (int i = 0; i < comparators.length; i++) {
comparators[i].copy(comparatorSlot, parentDoc);
}
// System.out.println(" startup: new OG doc=" + (docBase+parentDoc));
if (!trackMaxScore && trackScores) {
score = scorer.score();
}
final OneGroup og =
new OneGroup(comparatorSlot, docBase + parentDoc, score, joinScorers.length, trackScores);
og.readerContext = currentReaderContext;
copyGroups(og);
bottom = queue.add(og);
queueFull = totalHitCount == numParentHits;
if (queueFull) {
// End of startup transient: queue just filled up:
for (int i = 0; i < comparators.length; i++) {
comparators[i].setBottom(bottom.slot);
}
}
}
}
@Override
public LeafCollector getLeafCollector(final LeafReaderContext context) throws IOException {
final LeafFieldComparator[] comparators = queue.getComparators(context);
final int[] reverseMul = queue.getReverseMul();
final int docBase = context.docBase;
return new LeafCollector() {
private Scorer scorer;
@Override
public void setScorer(Scorer scorer) throws IOException {
// System.out.println("C.setScorer scorer=" + scorer);
// Since we invoke .score(), and the comparators likely
// do as well, cache it so it's only "really" computed
// once:
if (scorer instanceof ScoreCachingWrappingScorer == false) {
scorer = new ScoreCachingWrappingScorer(scorer);
}
this.scorer = scorer;
for (LeafFieldComparator comparator : comparators) {
comparator.setScorer(scorer);
}
Arrays.fill(joinScorers, null);
Queue<Scorer> queue = new LinkedList<>();
// System.out.println("\nqueue: add top scorer=" + scorer);
queue.add(scorer);
while ((scorer = queue.poll()) != null) {
// System.out.println(" poll: " + scorer + "; " + scorer.getWeight().getQuery());
if (scorer instanceof ToParentBlockJoinQuery.BlockJoinScorer) {
enroll(
(ToParentBlockJoinQuery) scorer.getWeight().getQuery(),
(ToParentBlockJoinQuery.BlockJoinScorer) scorer);
}
for (ChildScorer sub : scorer.getChildren()) {
// System.out.println(" add sub: " + sub.child + "; " +
// sub.child.getWeight().getQuery());
queue.add(sub.child);
}
}
}
@Override
public void collect(int parentDoc) throws IOException {
// System.out.println("\nC parentDoc=" + parentDoc);
totalHitCount++;
float score = Float.NaN;
if (trackMaxScore) {
score = scorer.score();
maxScore = Math.max(maxScore, score);
}
// TODO: we could sweep all joinScorers here and
// aggregate total child hit count, so we can fill this
// in getTopGroups (we wire it to 0 now)
if (queueFull) {
// System.out.println(" queueFull");
// Fastmatch: return if this hit is not competitive
int c = 0;
for (int i = 0; i < comparators.length; ++i) {
c = reverseMul[i] * comparators[i].compareBottom(parentDoc);
if (c != 0) {
break;
}
}
if (c <= 0) { // in case of equality, this hit is not competitive as docs are visited in
// order
// Definitely not competitive.
// System.out.println(" skip");
return;
}
// System.out.println(" competes! doc=" + (docBase + parentDoc));
// This hit is competitive - replace bottom element in queue & adjustTop
for (LeafFieldComparator comparator : comparators) {
comparator.copy(bottom.slot, parentDoc);
}
if (!trackMaxScore && trackScores) {
score = scorer.score();
}
bottom.doc = docBase + parentDoc;
bottom.readerContext = context;
bottom.score = score;
copyGroups(bottom);
bottom = queue.updateTop();
for (LeafFieldComparator comparator : comparators) {
comparator.setBottom(bottom.slot);
}
} else {
// Startup transient: queue is not yet full:
final int comparatorSlot = totalHitCount - 1;
// Copy hit into queue
for (LeafFieldComparator comparator : comparators) {
comparator.copy(comparatorSlot, parentDoc);
}
// System.out.println(" startup: new OG doc=" + (docBase+parentDoc));
if (!trackMaxScore && trackScores) {
score = scorer.score();
}
final OneGroup og =
new OneGroup(
comparatorSlot, docBase + parentDoc, score, joinScorers.length, trackScores);
og.readerContext = context;
copyGroups(og);
bottom = queue.add(og);
queueFull = totalHitCount == numParentHits;
if (queueFull) {
// End of startup transient: queue just filled up:
for (LeafFieldComparator comparator : comparators) {
comparator.setBottom(bottom.slot);
}
}
}
}
// Pulls out child doc and scores for all join queries:
private void copyGroups(OneGroup og) {
// While rare, it's possible top arrays could be too
// short if join query had null scorer on first
// segment(s) but then became non-null on later segments
final int numSubScorers = joinScorers.length;
if (og.docs.length < numSubScorers) {
// While rare, this could happen if join query had
// null scorer on first segment(s) but then became
// non-null on later segments
og.docs = ArrayUtil.grow(og.docs, numSubScorers);
}
if (og.counts.length < numSubScorers) {
og.counts = ArrayUtil.grow(og.counts);
}
if (trackScores && og.scores.length < numSubScorers) {
og.scores = ArrayUtil.grow(og.scores, numSubScorers);
}
// System.out.println("\ncopyGroups parentDoc=" + og.doc);
for (int scorerIDX = 0; scorerIDX < numSubScorers; scorerIDX++) {
final ToParentBlockJoinQuery.BlockJoinScorer joinScorer = joinScorers[scorerIDX];
// System.out.println(" scorer=" + joinScorer);
if (joinScorer != null && docBase + joinScorer.getParentDoc() == og.doc) {
og.counts[scorerIDX] = joinScorer.getChildCount();
// System.out.println(" count=" + og.counts[scorerIDX]);
og.docs[scorerIDX] = joinScorer.swapChildDocs(og.docs[scorerIDX]);
assert og.docs[scorerIDX].length >= og.counts[scorerIDX]
: "length=" + og.docs[scorerIDX].length + " vs count=" + og.counts[scorerIDX];
// System.out.println(" len=" + og.docs[scorerIDX].length);
/*
for(int idx=0;idx<og.counts[scorerIDX];idx++) {
System.out.println(" docs[" + idx + "]=" + og.docs[scorerIDX][idx]);
}
*/
if (trackScores) {
// System.out.println(" copy scores");
og.scores[scorerIDX] = joinScorer.swapChildScores(og.scores[scorerIDX]);
assert og.scores[scorerIDX].length >= og.counts[scorerIDX]
: "length=" + og.scores[scorerIDX].length + " vs count=" + og.counts[scorerIDX];
}
} else {
og.counts[scorerIDX] = 0;
}
}
}
};
}
