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

Example 16 with LeafBucketCollectorBase

use of org.elasticsearch.search.aggregations.LeafBucketCollectorBase in project elasticsearch by elastic.

the class TopHitsAggregator method getLeafCollector.

@Override
public LeafBucketCollector getLeafCollector(final LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
    for (LongObjectPagedHashMap.Cursor<TopDocsAndLeafCollector> cursor : topDocsCollectors) {
        cursor.value.leafCollector = cursor.value.topLevelCollector.getLeafCollector(ctx);
    }
    return new LeafBucketCollectorBase(sub, null) {

        Scorer scorer;

        @Override
        public void setScorer(Scorer scorer) throws IOException {
            this.scorer = scorer;
            for (LongObjectPagedHashMap.Cursor<TopDocsAndLeafCollector> cursor : topDocsCollectors) {
                cursor.value.leafCollector.setScorer(scorer);
            }
            super.setScorer(scorer);
        }

        @Override
        public void collect(int docId, long bucket) throws IOException {
            TopDocsAndLeafCollector collectors = topDocsCollectors.get(bucket);
            if (collectors == null) {
                SortAndFormats sort = subSearchContext.sort();
                int topN = subSearchContext.from() + subSearchContext.size();
                if (sort == null) {
                    for (RescoreSearchContext rescoreContext : context.rescore()) {
                        topN = Math.max(rescoreContext.window(), 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());
                TopDocsCollector<?> topLevelCollector;
                if (sort == null) {
                    topLevelCollector = TopScoreDocCollector.create(topN);
                } else {
                    topLevelCollector = TopFieldCollector.create(sort.sort, topN, true, subSearchContext.trackScores(), subSearchContext.trackScores());
                }
                collectors = new TopDocsAndLeafCollector(topLevelCollector);
                collectors.leafCollector = collectors.topLevelCollector.getLeafCollector(ctx);
                collectors.leafCollector.setScorer(scorer);
                topDocsCollectors.put(bucket, collectors);
            }
            collectors.leafCollector.collect(docId);
        }
    };
}

Example 17 with LeafBucketCollectorBase

use of org.elasticsearch.search.aggregations.LeafBucketCollectorBase in project elasticsearch by elastic.

the class ScriptedMetricAggregator method getLeafCollector.

@Override
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
    final LeafSearchScript leafMapScript = mapScript.getLeafSearchScript(ctx);
    return new LeafBucketCollectorBase(sub, mapScript) {

        @Override
        public void collect(int doc, long bucket) throws IOException {
            assert bucket == 0 : bucket;
            leafMapScript.setDocument(doc);
            leafMapScript.run();
        }
    };
}

Example 18 with LeafBucketCollectorBase

use of org.elasticsearch.search.aggregations.LeafBucketCollectorBase in project elasticsearch by elastic.

the class StatsAggregator method getLeafCollector.

@Override
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final 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 {
            if (bucket >= counts.size()) {
                final long from = counts.size();
                final long overSize = BigArrays.overSize(bucket + 1);
                counts = bigArrays.resize(counts, overSize);
                sums = bigArrays.resize(sums, overSize);
                mins = bigArrays.resize(mins, overSize);
                maxes = bigArrays.resize(maxes, overSize);
                mins.fill(from, overSize, Double.POSITIVE_INFINITY);
                maxes.fill(from, overSize, Double.NEGATIVE_INFINITY);
            }
            values.setDocument(doc);
            final int valuesCount = values.count();
            counts.increment(bucket, valuesCount);
            double sum = 0;
            double min = mins.get(bucket);
            double max = maxes.get(bucket);
            for (int i = 0; i < valuesCount; i++) {
                double value = values.valueAt(i);
                sum += value;
                min = Math.min(min, value);
                max = Math.max(max, value);
            }
            sums.increment(bucket, sum);
            mins.set(bucket, min);
            maxes.set(bucket, max);
        }
    };
}

Example 19 with LeafBucketCollectorBase

use of org.elasticsearch.search.aggregations.LeafBucketCollectorBase in project elasticsearch by elastic.

the class AvgAggregator method getLeafCollector.

@Override
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final 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 {
            counts = bigArrays.grow(counts, bucket + 1);
            sums = bigArrays.grow(sums, bucket + 1);
            values.setDocument(doc);
            final int valueCount = values.count();
            counts.increment(bucket, valueCount);
            double sum = 0;
            for (int i = 0; i < valueCount; i++) {
                sum += values.valueAt(i);
            }
            sums.increment(bucket, sum);
        }
    };
}

Example 20 with LeafBucketCollectorBase

use of org.elasticsearch.search.aggregations.LeafBucketCollectorBase in project elasticsearch by elastic.

the class MaxAggregator method getLeafCollector.

@Override
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
    if (valuesSource == null) {
        return LeafBucketCollector.NO_OP_COLLECTOR;
    }
    final BigArrays bigArrays = context.bigArrays();
    final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx);
    final NumericDoubleValues values = MultiValueMode.MAX.select(allValues, Double.NEGATIVE_INFINITY);
    return new LeafBucketCollectorBase(sub, allValues) {

        @Override
        public void collect(int doc, long bucket) throws IOException {
            if (bucket >= maxes.size()) {
                long from = maxes.size();
                maxes = bigArrays.grow(maxes, bucket + 1);
                maxes.fill(from, maxes.size(), Double.NEGATIVE_INFINITY);
            }
            final double value = values.get(doc);
            double max = maxes.get(bucket);
            max = Math.max(max, value);
            maxes.set(bucket, max);
        }
    };
}