Fix TopHitsAggregator to not ignore the top-level/leaf collector split.

This commit is contained in:
Adrien Grand 2014-11-04 11:32:24 +01:00
parent 042fadc860
commit 7f10129364

View file

@ -40,9 +40,19 @@ import java.io.IOException;
*/ */
public class TopHitsAggregator extends MetricsAggregator implements ScorerAware { public class TopHitsAggregator extends MetricsAggregator implements ScorerAware {
/** Simple wrapper around a top-level collector and the current leaf collector. */
private static class TopDocsAndLeafCollector {
final TopDocsCollector<?> topLevelCollector;
LeafCollector leafCollector;
TopDocsAndLeafCollector(TopDocsCollector<?> topLevelCollector) {
this.topLevelCollector = topLevelCollector;
}
}
private final FetchPhase fetchPhase; private final FetchPhase fetchPhase;
private final TopHitsContext topHitsContext; private final TopHitsContext topHitsContext;
private final LongObjectPagedHashMap<TopDocsCollector> topDocsCollectors; private final LongObjectPagedHashMap<TopDocsAndLeafCollector> topDocsCollectors;
private Scorer currentScorer; private Scorer currentScorer;
private LeafReaderContext currentContext; private LeafReaderContext currentContext;
@ -62,11 +72,11 @@ public class TopHitsAggregator extends MetricsAggregator implements ScorerAware
@Override @Override
public InternalAggregation buildAggregation(long owningBucketOrdinal) { public InternalAggregation buildAggregation(long owningBucketOrdinal) {
TopDocsCollector topDocsCollector = topDocsCollectors.get(owningBucketOrdinal); TopDocsAndLeafCollector topDocsCollector = topDocsCollectors.get(owningBucketOrdinal);
if (topDocsCollector == null) { if (topDocsCollector == null) {
return buildEmptyAggregation(); return buildEmptyAggregation();
} else { } else {
TopDocs topDocs = topDocsCollector.topDocs(); TopDocs topDocs = topDocsCollector.topLevelCollector.topDocs();
if (topDocs.totalHits == 0) { if (topDocs.totalHits == 0) {
return buildEmptyAggregation(); return buildEmptyAggregation();
} }
@ -101,32 +111,25 @@ public class TopHitsAggregator extends MetricsAggregator implements ScorerAware
@Override @Override
public void collect(int docId, long bucketOrdinal) throws IOException { public void collect(int docId, long bucketOrdinal) throws IOException {
TopDocsCollector topDocsCollector = topDocsCollectors.get(bucketOrdinal); TopDocsAndLeafCollector collectors = topDocsCollectors.get(bucketOrdinal);
if (topDocsCollector == null) { if (collectors == null) {
Sort sort = topHitsContext.sort(); Sort sort = topHitsContext.sort();
int topN = topHitsContext.from() + topHitsContext.size(); int topN = topHitsContext.from() + topHitsContext.size();
topDocsCollectors.put( TopDocsCollector<?> topLevelCollector = sort != null ? TopFieldCollector.create(sort, topN, true, topHitsContext.trackScores(), topHitsContext.trackScores(), false) : TopScoreDocCollector.create(topN, false);
bucketOrdinal, collectors = new TopDocsAndLeafCollector(topLevelCollector);
topDocsCollector = sort != null ? TopFieldCollector.create(sort, topN, true, topHitsContext.trackScores(), topHitsContext.trackScores(), false) : TopScoreDocCollector.create(topN, false) collectors.leafCollector = collectors.topLevelCollector.getLeafCollector(currentContext);
); collectors.leafCollector.setScorer(currentScorer);
// TODO: this is bogus, it only works because TopDocsCollector subclasses SimpleCollector, topDocsCollectors.put(bucketOrdinal, collectors);
// we should not be ignoring the return value: instead make it work properly per-segment
LeafCollector ignoredReturnValue = topDocsCollector.getLeafCollector(currentContext);
assert ignoredReturnValue == topDocsCollector;
topDocsCollector.setScorer(currentScorer);
} }
topDocsCollector.collect(docId); collectors.leafCollector.collect(docId);
} }
@Override @Override
public void setNextReader(LeafReaderContext context) { public void setNextReader(LeafReaderContext context) {
this.currentContext = context; this.currentContext = context;
for (LongObjectPagedHashMap.Cursor<TopDocsCollector> cursor : topDocsCollectors) { for (LongObjectPagedHashMap.Cursor<TopDocsAndLeafCollector> cursor : topDocsCollectors) {
try { try {
// TODO: this is bogus, it only works because TopDocsCollector subclasses SimpleCollector, cursor.value.leafCollector = cursor.value.topLevelCollector.getLeafCollector(currentContext);
// we should not be ignoring the return value: instead make it work properly per-segment
LeafCollector ignoredReturnValue = cursor.value.getLeafCollector(context);
assert ignoredReturnValue == cursor.value;
} catch (IOException e) { } catch (IOException e) {
throw ExceptionsHelper.convertToElastic(e); throw ExceptionsHelper.convertToElastic(e);
} }
@ -136,9 +139,9 @@ public class TopHitsAggregator extends MetricsAggregator implements ScorerAware
@Override @Override
public void setScorer(Scorer scorer) { public void setScorer(Scorer scorer) {
this.currentScorer = scorer; this.currentScorer = scorer;
for (LongObjectPagedHashMap.Cursor<TopDocsCollector> cursor : topDocsCollectors) { for (LongObjectPagedHashMap.Cursor<TopDocsAndLeafCollector> cursor : topDocsCollectors) {
try { try {
cursor.value.setScorer(scorer); cursor.value.leafCollector.setScorer(scorer);
} catch (IOException e) { } catch (IOException e) {
throw ExceptionsHelper.convertToElastic(e); throw ExceptionsHelper.convertToElastic(e);
} }