Handle TopFieldDocs copy in TopDocsCollectorContext

This commit fixes the clone of TopFieldDocs.

Relates #37179
Relates #37266
This commit is contained in:
Jim Ferenczi 2019-01-10 00:26:35 +01:00
parent 234059d2c0
commit 586093ec5e
1 changed files with 28 additions and 17 deletions

View File

@ -35,6 +35,7 @@ import org.apache.lucene.search.TermQuery;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.TopDocsCollector;
import org.apache.lucene.search.TopFieldCollector;
import org.apache.lucene.search.TopFieldDocs;
import org.apache.lucene.search.TopScoreDocCollector;
import org.apache.lucene.search.TotalHitCountCollector;
import org.apache.lucene.search.TotalHits;
@ -189,11 +190,11 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
}
}
private final Collector collector;
protected final @Nullable SortAndFormats sortAndFormats;
protected final Supplier<TotalHits> totalHitsSupplier;
protected final Supplier<TopDocs> topDocsSupplier;
protected final Supplier<Float> maxScoreSupplier;
private final Collector collector;
private final Supplier<TotalHits> totalHitsSupplier;
private final Supplier<TopDocs> topDocsSupplier;
private final Supplier<Float> maxScoreSupplier;
/**
* Ctr
@ -262,11 +263,23 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
return collector;
}
TopDocsAndMaxScore newTopDocs() {
TopDocs in = topDocsSupplier.get();
float maxScore = maxScoreSupplier.get();
final TopDocs newTopDocs;
if (in instanceof TopFieldDocs) {
TopFieldDocs fieldDocs = (TopFieldDocs) in;
newTopDocs = new TopFieldDocs(totalHitsSupplier.get(), fieldDocs.scoreDocs, fieldDocs.fields);
} else {
newTopDocs = new TopDocs(totalHitsSupplier.get(), in.scoreDocs);
}
return new TopDocsAndMaxScore(newTopDocs, maxScore);
}
@Override
void postProcess(QuerySearchResult result) throws IOException {
final TopDocs topDocs = new TopDocs(totalHitsSupplier.get(), topDocsSupplier.get().scoreDocs);
result.topDocs(new TopDocsAndMaxScore(topDocs, maxScoreSupplier.get()),
sortAndFormats == null ? null : sortAndFormats.formats);
final TopDocsAndMaxScore topDocs = newTopDocs();
result.topDocs(topDocs, sortAndFormats == null ? null : sortAndFormats.formats);
}
}
@ -291,27 +304,25 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
@Override
void postProcess(QuerySearchResult result) throws IOException {
final TopDocs topDocs = new TopDocs(totalHitsSupplier.get(), topDocsSupplier.get().scoreDocs);
final float maxScore;
final TopDocsAndMaxScore topDocs = newTopDocs();
if (scrollContext.totalHits == null) {
// first round
scrollContext.totalHits = topDocs.totalHits;
maxScore = scrollContext.maxScore = maxScoreSupplier.get();
scrollContext.totalHits = topDocs.topDocs.totalHits;
scrollContext.maxScore = topDocs.maxScore;
} else {
// subsequent round: the total number of hits and
// the maximum score were computed on the first round
topDocs.totalHits = scrollContext.totalHits;
maxScore = scrollContext.maxScore;
topDocs.topDocs.totalHits = scrollContext.totalHits;
topDocs.maxScore = scrollContext.maxScore;
}
if (numberOfShards == 1) {
// if we fetch the document in the same roundtrip, we already know the last emitted doc
if (topDocs.scoreDocs.length > 0) {
if (topDocs.topDocs.scoreDocs.length > 0) {
// set the last emitted doc
scrollContext.lastEmittedDoc = topDocs.scoreDocs[topDocs.scoreDocs.length - 1];
scrollContext.lastEmittedDoc = topDocs.topDocs.scoreDocs[topDocs.topDocs.scoreDocs.length - 1];
}
}
result.topDocs(new TopDocsAndMaxScore(topDocs, maxScore),
sortAndFormats == null ? null : sortAndFormats.formats);
result.topDocs(topDocs, sortAndFormats == null ? null : sortAndFormats.formats);
}
}