Examples with LeafBucketCollectorBase org.opensearch.search.aggregations.LeafBucketCollectorBase used on opensource projects

Example 26 with LeafBucketCollectorBase

use of org.opensearch.search.aggregations.LeafBucketCollectorBase in project OpenSearch by opensearch-project.

the class TopHitsAggregator method getLeafCollector.

@Override
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
    // Create leaf collectors here instead of at the aggregator level. Otherwise in case this collector get invoked
    // 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<LeafCollector> leafCollectors = new LongObjectHashMap<>(1);
    return new LeafBucketCollectorBase(sub, null) {

        Scorable scorer;

        @Override
        public void setScorer(Scorable scorer) throws IOException {
            this.scorer = scorer;
            super.setScorer(scorer);
            for (ObjectCursor<LeafCollector> cursor : leafCollectors.values()) {
                cursor.value.setScorer(scorer);
            }
        }

        @Override
        public void collect(int docId, long bucket) throws IOException {
            Collectors collectors = topDocsCollectors.get(bucket);
            if (collectors == null) {
                SortAndFormats sort = subSearchContext.sort();
                int topN = subSearchContext.from() + subSearchContext.size();
                if (sort == null) {
                    for (RescoreContext rescoreContext : context.rescore()) {
                        topN = Math.max(rescoreContext.getWindowSize(), topN);
                    }
                }
                // In the QueryPhase we don't need this protection, because it is build into the IndexSearcher,
                // but here we create collectors ourselves and we need prevent OOM because of crazy an offset and size.
                topN = Math.min(topN, subSearchContext.searcher().getIndexReader().maxDoc());
                if (sort == null) {
                    collectors = new Collectors(TopScoreDocCollector.create(topN, Integer.MAX_VALUE), null);
                } else {
                    // TODO: can we pass trackTotalHits=subSearchContext.trackTotalHits(){
                    // Note that this would require to catch CollectionTerminatedException
                    collectors = new Collectors(TopFieldCollector.create(sort.sort, topN, Integer.MAX_VALUE), subSearchContext.trackScores() ? new MaxScoreCollector() : null);
                }
                topDocsCollectors.put(bucket, collectors);
            }
            final LeafCollector leafCollector;
            final int key = leafCollectors.indexOf(bucket);
            if (key < 0) {
                leafCollector = collectors.collector.getLeafCollector(ctx);
                if (scorer != null) {
                    leafCollector.setScorer(scorer);
                }
                leafCollectors.indexInsert(key, bucket, leafCollector);
            } else {
                leafCollector = leafCollectors.indexGet(key);
            }
            leafCollector.collect(docId);
        }
    };
}

Example 27 with LeafBucketCollectorBase

use of org.opensearch.search.aggregations.LeafBucketCollectorBase in project OpenSearch by opensearch-project.

the class MedianAbsoluteDeviationAggregator method getLeafCollector.

@Override
protected LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
    if (valuesSource == null) {
        return LeafBucketCollector.NO_OP_COLLECTOR;
    }
    final BigArrays bigArrays = context.bigArrays();
    final SortedNumericDoubleValues values = valuesSource.doubleValues(ctx);
    return new LeafBucketCollectorBase(sub, values) {

        @Override
        public void collect(int doc, long bucket) throws IOException {
            valueSketches = bigArrays.grow(valueSketches, bucket + 1);
            TDigestState valueSketch = valueSketches.get(bucket);
            if (valueSketch == null) {
                valueSketch = new TDigestState(compression);
                valueSketches.set(bucket, valueSketch);
            }
            if (values.advanceExact(doc)) {
                final int valueCount = values.docValueCount();
                for (int i = 0; i < valueCount; i++) {
                    final double value = values.nextValue();
                    valueSketch.add(value);
                }
            }
        }
    };
}