diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/TopHitsAggregator.java b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/TopHitsAggregator.java index 7b38bf2d2f25b..bd47c5b35c380 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/TopHitsAggregator.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/TopHitsAggregator.java @@ -8,9 +8,6 @@ package org.elasticsearch.search.aggregations.metrics; -import com.carrotsearch.hppc.LongObjectHashMap; -import com.carrotsearch.hppc.cursors.ObjectCursor; - import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.search.Collector; import org.apache.lucene.search.FieldDoc; @@ -29,7 +26,9 @@ import org.elasticsearch.action.search.MaxScoreCollector; import org.elasticsearch.common.lucene.Lucene; import org.elasticsearch.common.lucene.search.TopDocsAndMaxScore; +import org.elasticsearch.common.util.BigArrays; import org.elasticsearch.common.util.LongObjectPagedHashMap; +import org.elasticsearch.common.util.LongObjectPagedHashMap.Cursor; import org.elasticsearch.core.Releasables; import org.elasticsearch.search.SearchHit; import org.elasticsearch.search.SearchHits; @@ -65,9 +64,12 @@ private static class Collectors { } } + private final BigArrays bigArrays; private final SubSearchContext subSearchContext; private final LongObjectPagedHashMap topDocsCollectors; private final List fetchProfiles; + // this must be mutable so it can be closed/replaced on each call to getLeafCollector + private LongObjectPagedHashMap leafCollectors; TopHitsAggregator( SubSearchContext subSearchContext, @@ -77,9 +79,10 @@ private static class Collectors { Map metadata ) throws IOException { super(name, context, parent, metadata); - topDocsCollectors = new LongObjectPagedHashMap<>(1, context.bigArrays()); + this.bigArrays = context.bigArrays(); this.subSearchContext = subSearchContext; - fetchProfiles = context.profiling() ? new ArrayList<>() : null; + this.topDocsCollectors = new LongObjectPagedHashMap<>(1, bigArrays); + this.fetchProfiles = context.profiling() ? new ArrayList<>() : null; } @Override @@ -99,7 +102,11 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCol // when post collecting then we have already replaced the leaf readers on the aggregator level have already been // replaced with the next leaf readers and then post collection pushes docids of the previous segment, which // then causes assertions to trip or incorrect top docs to be computed. - final LongObjectHashMap leafCollectors = new LongObjectHashMap<>(1); + if (leafCollectors != null) { + leafCollectors.close(); + leafCollectors = null; // set to null, just in case the new allocation below fails + } + leafCollectors = new LongObjectPagedHashMap<>(1, bigArrays); return new LeafBucketCollectorBase(sub, null) { Scorable scorer; @@ -108,8 +115,8 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCol public void setScorer(Scorable scorer) throws IOException { this.scorer = scorer; super.setScorer(scorer); - for (ObjectCursor cursor : leafCollectors.values()) { - cursor.value.setScorer(scorer); + for (Cursor leafCollector : leafCollectors) { + leafCollector.value.setScorer(scorer); } } @@ -140,16 +147,13 @@ public void collect(int docId, long bucket) throws IOException { topDocsCollectors.put(bucket, collectors); } - final LeafCollector leafCollector; - final int key = leafCollectors.indexOf(bucket); - if (key < 0) { + LeafCollector leafCollector = leafCollectors.get(bucket); + if (leafCollector == null) { leafCollector = collectors.collector.getLeafCollector(ctx); if (scorer != null) { leafCollector.setScorer(scorer); } - leafCollectors.indexInsert(key, bucket, leafCollector); - } else { - leafCollector = leafCollectors.indexGet(key); + leafCollectors.put(bucket, leafCollector); } leafCollector.collect(docId); } @@ -250,6 +254,6 @@ public void collectDebugInfo(BiConsumer add) { @Override protected void doClose() { - Releasables.close(topDocsCollectors); + Releasables.close(topDocsCollectors, leafCollectors); } }