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

Example 21 with LeafBucketCollectorBase

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

the class ReverseNestedAggregator method getLeafCollector.

@Override
protected LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
    // In ES if parent is deleted, then also the children are deleted, so the child docs this agg receives
    // must belong to parent docs that is alive. For this reason acceptedDocs can be null here.
    final BitSet parentDocs = parentBitsetProducer.getBitSet(ctx);
    if (parentDocs == null) {
        return LeafBucketCollector.NO_OP_COLLECTOR;
    }
    final LongIntHashMap bucketOrdToLastCollectedParentDoc = new LongIntHashMap(32);
    return new LeafBucketCollectorBase(sub, null) {

        @Override
        public void collect(int childDoc, long bucket) throws IOException {
            // fast forward to retrieve the parentDoc this childDoc belongs to
            final int parentDoc = parentDocs.nextSetBit(childDoc);
            assert childDoc <= parentDoc && parentDoc != DocIdSetIterator.NO_MORE_DOCS;
            int keySlot = bucketOrdToLastCollectedParentDoc.indexOf(bucket);
            if (bucketOrdToLastCollectedParentDoc.indexExists(keySlot)) {
                int lastCollectedParentDoc = bucketOrdToLastCollectedParentDoc.indexGet(keySlot);
                if (parentDoc > lastCollectedParentDoc) {
                    collectBucket(sub, parentDoc, bucket);
                    bucketOrdToLastCollectedParentDoc.indexReplace(keySlot, parentDoc);
                }
            } else {
                collectBucket(sub, parentDoc, bucket);
                bucketOrdToLastCollectedParentDoc.indexInsert(keySlot, bucket, parentDoc);
            }
        }
    };
}

Example 22 with LeafBucketCollectorBase

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

the class RangeAggregator method getLeafCollector.

@Override
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
    final SortedNumericDoubleValues values = valuesSource.doubleValues(ctx);
    return new LeafBucketCollectorBase(sub, values) {

        @Override
        public void collect(int doc, long bucket) throws IOException {
            values.setDocument(doc);
            final int valuesCount = values.count();
            for (int i = 0, lo = 0; i < valuesCount; ++i) {
                final double value = values.valueAt(i);
                lo = collect(doc, value, bucket, lo);
            }
        }

        private int collect(int doc, double value, long owningBucketOrdinal, int lowBound) throws IOException {
            // all candidates are between these indexes
            int lo = lowBound, hi = ranges.length - 1;
            int mid = (lo + hi) >>> 1;
            while (lo <= hi) {
                if (value < ranges[mid].from) {
                    hi = mid - 1;
                } else if (value >= maxTo[mid]) {
                    lo = mid + 1;
                } else {
                    break;
                }
                mid = (lo + hi) >>> 1;
            }
            // no potential candidate
            if (lo > hi)
                return lo;
            // binary search the lower bound
            int startLo = lo, startHi = mid;
            while (startLo <= startHi) {
                final int startMid = (startLo + startHi) >>> 1;
                if (value >= maxTo[startMid]) {
                    startLo = startMid + 1;
                } else {
                    startHi = startMid - 1;
                }
            }
            // binary search the upper bound
            int endLo = mid, endHi = hi;
            while (endLo <= endHi) {
                final int endMid = (endLo + endHi) >>> 1;
                if (value < ranges[endMid].from) {
                    endHi = endMid - 1;
                } else {
                    endLo = endMid + 1;
                }
            }
            assert startLo == lowBound || value >= maxTo[startLo - 1];
            assert endHi == ranges.length - 1 || value < ranges[endHi + 1].from;
            for (int i = startLo; i <= endHi; ++i) {
                if (ranges[i].matches(value)) {
                    collectBucket(sub, doc, subBucketOrdinal(owningBucketOrdinal, i));
                }
            }
            return endHi + 1;
        }
    };
}

Example 23 with LeafBucketCollectorBase

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

the class GeoHashGridAggregator method getLeafCollector.

@Override
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
    final SortedNumericDocValues values = valuesSource.longValues(ctx);
    return new LeafBucketCollectorBase(sub, null) {

        @Override
        public void collect(int doc, long bucket) throws IOException {
            assert bucket == 0;
            values.setDocument(doc);
            final int valuesCount = values.count();
            long previous = Long.MAX_VALUE;
            for (int i = 0; i < valuesCount; ++i) {
                final long val = values.valueAt(i);
                if (previous != val || i == 0) {
                    long bucketOrdinal = bucketOrds.add(val);
                    if (bucketOrdinal < 0) {
                        // already seen
                        bucketOrdinal = -1 - bucketOrdinal;
                        collectExistingBucket(sub, doc, bucketOrdinal);
                    } else {
                        collectBucket(sub, doc, bucketOrdinal);
                    }
                    previous = val;
                }
            }
        }
    };
}