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

Example 16 with LeafBucketCollectorBase

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

the class NestedAggregator method getLeafCollector.

@Override
public LeafBucketCollector getLeafCollector(final LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
    IndexReaderContext topLevelContext = ReaderUtil.getTopLevelContext(ctx);
    IndexSearcher searcher = new IndexSearcher(topLevelContext);
    searcher.setQueryCache(null);
    Weight weight = searcher.createWeight(searcher.rewrite(childFilter), ScoreMode.COMPLETE_NO_SCORES, 1f);
    Scorer childDocsScorer = weight.scorer(ctx);
    final BitSet parentDocs = parentFilter.getBitSet(ctx);
    final DocIdSetIterator childDocs = childDocsScorer != null ? childDocsScorer.iterator() : null;
    if (collectsFromSingleBucket) {
        return new LeafBucketCollectorBase(sub, null) {

            @Override
            public void collect(int parentDoc, long bucket) throws IOException {
                // doc), so we can skip:
                if (parentDoc == 0 || parentDocs == null || childDocs == null) {
                    return;
                }
                final int prevParentDoc = parentDocs.prevSetBit(parentDoc - 1);
                int childDocId = childDocs.docID();
                if (childDocId <= prevParentDoc) {
                    childDocId = childDocs.advance(prevParentDoc + 1);
                }
                for (; childDocId < parentDoc; childDocId = childDocs.nextDoc()) {
                    collectBucket(sub, childDocId, bucket);
                }
            }
        };
    } else {
        return bufferingNestedLeafBucketCollector = new BufferingNestedLeafBucketCollector(sub, parentDocs, childDocs);
    }
}

Example 17 with LeafBucketCollectorBase

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

the class DateRangeHistogramAggregator method getLeafCollector.

@Override
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
    if (valuesSource == null) {
        return LeafBucketCollector.NO_OP_COLLECTOR;
    }
    SortedBinaryDocValues values = valuesSource.bytesValues(ctx);
    RangeType rangeType = valuesSource.rangeType();
    return new LeafBucketCollectorBase(sub, values) {

        @Override
        public void collect(int doc, long owningBucketOrd) throws IOException {
            if (values.advanceExact(doc)) {
                // Is it possible for valuesCount to be > 1 here? Multiple ranges are encoded into the same BytesRef in the binary doc
                // values, so it isn't clear what we'd be iterating over.
                int valuesCount = values.docValueCount();
                assert valuesCount == 1 : "Value count for ranges should always be 1";
                long previousKey = Long.MIN_VALUE;
                for (int i = 0; i < valuesCount; i++) {
                    BytesRef encodedRanges = values.nextValue();
                    List<RangeFieldMapper.Range> ranges = rangeType.decodeRanges(encodedRanges);
                    long previousFrom = Long.MIN_VALUE;
                    for (RangeFieldMapper.Range range : ranges) {
                        Long from = (Long) range.getFrom();
                        // The encoding should ensure that this assert is always true.
                        assert from >= previousFrom : "Start of range not >= previous start";
                        final Long to = (Long) range.getTo();
                        final long effectiveFrom = (hardBounds != null && hardBounds.getMin() != null) ? max(from, hardBounds.getMin()) : from;
                        final long effectiveTo = (hardBounds != null && hardBounds.getMax() != null) ? min(to, hardBounds.getMax()) : to;
                        final long startKey = preparedRounding.round(effectiveFrom);
                        final long endKey = preparedRounding.round(effectiveTo);
                        for (long key = max(startKey, previousKey); key <= endKey; key = preparedRounding.nextRoundingValue(key)) {
                            if (key == previousKey) {
                                continue;
                            }
                            // Bucket collection identical to NumericHistogramAggregator, could be refactored
                            long bucketOrd = bucketOrds.add(owningBucketOrd, key);
                            if (bucketOrd < 0) {
                                // already seen
                                bucketOrd = -1 - bucketOrd;
                                collectExistingBucket(sub, doc, bucketOrd);
                            } else {
                                collectBucket(sub, doc, bucketOrd);
                            }
                        }
                        if (endKey > previousKey) {
                            previousKey = endKey;
                        }
                    }
                }
            }
        }
    };
}

Example 18 with LeafBucketCollectorBase

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

the class StringRareTermsAggregator method getLeafCollector.

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

        final BytesRefBuilder previous = new BytesRefBuilder();

        @Override
        public void collect(int docId, long owningBucketOrd) throws IOException {
            if (false == values.advanceExact(docId)) {
                return;
            }
            int valuesCount = values.docValueCount();
            previous.clear();
            // need to take care of dups
            for (int i = 0; i < valuesCount; ++i) {
                BytesRef bytes = values.nextValue();
                if (filter != null && false == filter.accept(bytes)) {
                    continue;
                }
                if (i > 0 && previous.get().equals(bytes)) {
                    continue;
                }
                previous.copyBytes(bytes);
                long bucketOrdinal = bucketOrds.add(owningBucketOrd, bytes);
                if (bucketOrdinal < 0) {
                    // already seen
                    bucketOrdinal = -1 - bucketOrdinal;
                    collectExistingBucket(sub, docId, bucketOrdinal);
                } else {
                    collectBucket(sub, docId, bucketOrdinal);
                }
            }
        }
    };
}

Example 19 with LeafBucketCollectorBase

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

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);
    final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
    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);
            compensations = bigArrays.grow(compensations, bucket + 1);
            if (values.advanceExact(doc)) {
                final int valueCount = values.docValueCount();
                counts.increment(bucket, valueCount);
                // Compute the sum of double values with Kahan summation algorithm which is more
                // accurate than naive summation.
                double sum = sums.get(bucket);
                double compensation = compensations.get(bucket);
                kahanSummation.reset(sum, compensation);
                for (int i = 0; i < valueCount; i++) {
                    double value = values.nextValue();
                    kahanSummation.add(value);
                }
                sums.set(bucket, kahanSummation.value());
                compensations.set(bucket, kahanSummation.delta());
            }
        }
    };
}

Example 20 with LeafBucketCollectorBase

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

the class GeoCentroidAggregator method getLeafCollector.

@Override
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
    if (valuesSource == null) {
        return LeafBucketCollector.NO_OP_COLLECTOR;
    }
    final BigArrays bigArrays = context.bigArrays();
    final MultiGeoPointValues values = valuesSource.geoPointValues(ctx);
    final CompensatedSum compensatedSumLat = new CompensatedSum(0, 0);
    final CompensatedSum compensatedSumLon = new CompensatedSum(0, 0);
    return new LeafBucketCollectorBase(sub, values) {

        @Override
        public void collect(int doc, long bucket) throws IOException {
            latSum = bigArrays.grow(latSum, bucket + 1);
            lonSum = bigArrays.grow(lonSum, bucket + 1);
            lonCompensations = bigArrays.grow(lonCompensations, bucket + 1);
            latCompensations = bigArrays.grow(latCompensations, bucket + 1);
            counts = bigArrays.grow(counts, bucket + 1);
            if (values.advanceExact(doc)) {
                final int valueCount = values.docValueCount();
                // increment by the number of points for this document
                counts.increment(bucket, valueCount);
                // Compute the sum of double values with Kahan summation algorithm which is more
                // accurate than naive summation.
                double sumLat = latSum.get(bucket);
                double compensationLat = latCompensations.get(bucket);
                double sumLon = lonSum.get(bucket);
                double compensationLon = lonCompensations.get(bucket);
                compensatedSumLat.reset(sumLat, compensationLat);
                compensatedSumLon.reset(sumLon, compensationLon);
                // update the sum
                for (int i = 0; i < valueCount; ++i) {
                    GeoPoint value = values.nextValue();
                    // latitude
                    compensatedSumLat.add(value.getLat());
                    // longitude
                    compensatedSumLon.add(value.getLon());
                }
                lonSum.set(bucket, compensatedSumLon.value());
                lonCompensations.set(bucket, compensatedSumLon.delta());
                latSum.set(bucket, compensatedSumLat.value());
                latCompensations.set(bucket, compensatedSumLat.delta());
            }
        }
    };
}