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