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

Example 11 with LeafBucketCollectorBase

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

the class ExtendedStatsAggregator 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 compensatedSum = new CompensatedSum(0, 0);
    final CompensatedSum compensatedSumOfSqr = new CompensatedSum(0, 0);
    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);
                compensations = bigArrays.resize(compensations, overSize);
                mins = bigArrays.resize(mins, overSize);
                maxes = bigArrays.resize(maxes, overSize);
                sumOfSqrs = bigArrays.resize(sumOfSqrs, overSize);
                compensationOfSqrs = bigArrays.resize(compensationOfSqrs, overSize);
                mins.fill(from, overSize, Double.POSITIVE_INFINITY);
                maxes.fill(from, overSize, Double.NEGATIVE_INFINITY);
            }
            if (values.advanceExact(doc)) {
                final int valuesCount = values.docValueCount();
                counts.increment(bucket, valuesCount);
                double min = mins.get(bucket);
                double max = maxes.get(bucket);
                // Compute the sum and sum of squires for double values with Kahan summation algorithm
                // which is more accurate than naive summation.
                double sum = sums.get(bucket);
                double compensation = compensations.get(bucket);
                compensatedSum.reset(sum, compensation);
                double sumOfSqr = sumOfSqrs.get(bucket);
                double compensationOfSqr = compensationOfSqrs.get(bucket);
                compensatedSumOfSqr.reset(sumOfSqr, compensationOfSqr);
                for (int i = 0; i < valuesCount; i++) {
                    double value = values.nextValue();
                    compensatedSum.add(value);
                    compensatedSumOfSqr.add(value * value);
                    min = Math.min(min, value);
                    max = Math.max(max, value);
                }
                sums.set(bucket, compensatedSum.value());
                compensations.set(bucket, compensatedSum.delta());
                sumOfSqrs.set(bucket, compensatedSumOfSqr.value());
                compensationOfSqrs.set(bucket, compensatedSumOfSqr.delta());
                mins.set(bucket, min);
                maxes.set(bucket, max);
            }
        }
    };
}

Example 12 with LeafBucketCollectorBase

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

the class GeoBoundsAggregator method getLeafCollector.

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

        @Override
        public void collect(int doc, long bucket) throws IOException {
            if (bucket >= tops.size()) {
                long from = tops.size();
                tops = bigArrays.grow(tops, bucket + 1);
                tops.fill(from, tops.size(), Double.NEGATIVE_INFINITY);
                bottoms = bigArrays.resize(bottoms, tops.size());
                bottoms.fill(from, bottoms.size(), Double.POSITIVE_INFINITY);
                posLefts = bigArrays.resize(posLefts, tops.size());
                posLefts.fill(from, posLefts.size(), Double.POSITIVE_INFINITY);
                posRights = bigArrays.resize(posRights, tops.size());
                posRights.fill(from, posRights.size(), Double.NEGATIVE_INFINITY);
                negLefts = bigArrays.resize(negLefts, tops.size());
                negLefts.fill(from, negLefts.size(), Double.POSITIVE_INFINITY);
                negRights = bigArrays.resize(negRights, tops.size());
                negRights.fill(from, negRights.size(), Double.NEGATIVE_INFINITY);
            }
            if (values.advanceExact(doc)) {
                final int valuesCount = values.docValueCount();
                for (int i = 0; i < valuesCount; ++i) {
                    GeoPoint value = values.nextValue();
                    double top = tops.get(bucket);
                    if (value.lat() > top) {
                        top = value.lat();
                    }
                    double bottom = bottoms.get(bucket);
                    if (value.lat() < bottom) {
                        bottom = value.lat();
                    }
                    double posLeft = posLefts.get(bucket);
                    if (value.lon() >= 0 && value.lon() < posLeft) {
                        posLeft = value.lon();
                    }
                    double posRight = posRights.get(bucket);
                    if (value.lon() >= 0 && value.lon() > posRight) {
                        posRight = value.lon();
                    }
                    double negLeft = negLefts.get(bucket);
                    if (value.lon() < 0 && value.lon() < negLeft) {
                        negLeft = value.lon();
                    }
                    double negRight = negRights.get(bucket);
                    if (value.lon() < 0 && value.lon() > negRight) {
                        negRight = value.lon();
                    }
                    tops.set(bucket, top);
                    bottoms.set(bucket, bottom);
                    posLefts.set(bucket, posLeft);
                    posRights.set(bucket, posRight);
                    negLefts.set(bucket, negLeft);
                    negRights.set(bucket, negRight);
                }
            }
        }
    };
}

Example 13 with LeafBucketCollectorBase

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

the class AbstractHDRPercentilesAggregator 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.Numeric) valuesSource).doubleValues(ctx);
    return new LeafBucketCollectorBase(sub, values) {

        @Override
        public void collect(int doc, long bucket) throws IOException {
            DoubleHistogram state = getExistingOrNewHistogram(bigArrays, bucket);
            if (values.advanceExact(doc)) {
                final int valueCount = values.docValueCount();
                for (int i = 0; i < valueCount; i++) {
                    state.recordValue(values.nextValue());
                }
            }
        }
    };
}

Example 14 with LeafBucketCollectorBase

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

the class DateHistogramAggregator method getLeafCollector.

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

        @Override
        public void collect(int doc, long owningBucketOrd) throws IOException {
            if (values.advanceExact(doc)) {
                int valuesCount = values.docValueCount();
                long previousRounded = Long.MIN_VALUE;
                for (int i = 0; i < valuesCount; ++i) {
                    long value = values.nextValue();
                    long rounded = preparedRounding.round(value);
                    assert rounded >= previousRounded;
                    if (rounded == previousRounded) {
                        continue;
                    }
                    if (hardBounds == null || hardBounds.contain(rounded)) {
                        long bucketOrd = bucketOrds.add(owningBucketOrd, rounded);
                        if (bucketOrd < 0) {
                            // already seen
                            bucketOrd = -1 - bucketOrd;
                            collectExistingBucket(sub, doc, bucketOrd);
                        } else {
                            collectBucket(sub, doc, bucketOrd);
                        }
                    }
                    previousRounded = rounded;
                }
            }
        }
    };
}

Example 15 with LeafBucketCollectorBase

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

the class RangeHistogramAggregator method getLeafCollector.

@Override
protected LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
    if (valuesSource == null) {
        return LeafBucketCollector.NO_OP_COLLECTOR;
    }
    final SortedBinaryDocValues values = valuesSource.bytesValues(ctx);
    final 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.
                final int valuesCount = values.docValueCount();
                assert valuesCount == 1 : "Value count for ranges should always be 1";
                double previousKey = Double.NEGATIVE_INFINITY;
                for (int i = 0; i < valuesCount; i++) {
                    BytesRef encodedRanges = values.nextValue();
                    List<RangeFieldMapper.Range> ranges = rangeType.decodeRanges(encodedRanges);
                    double previousFrom = Double.NEGATIVE_INFINITY;
                    for (RangeFieldMapper.Range range : ranges) {
                        final Double from = rangeType.doubleValue(range.getFrom());
                        // The encoding should ensure that this assert is always true.
                        assert from >= previousFrom : "Start of range not >= previous start";
                        final Double to = rangeType.doubleValue(range.getTo());
                        final double effectiveFrom = (hardBounds != null && hardBounds.getMin() != null) ? Double.max(from, hardBounds.getMin()) : from;
                        final double effectiveTo = (hardBounds != null && hardBounds.getMax() != null) ? Double.min(to, hardBounds.getMax()) : to;
                        final double startKey = Math.floor((effectiveFrom - offset) / interval);
                        final double endKey = Math.floor((effectiveTo - offset) / interval);
                        for (double key = Math.max(startKey, previousKey); key <= endKey; key++) {
                            if (key == previousKey) {
                                continue;
                            }
                            // Bucket collection identical to NumericHistogramAggregator, could be refactored
                            long bucketOrd = bucketOrds.add(owningBucketOrd, Double.doubleToLongBits(key));
                            if (bucketOrd < 0) {
                                // already seen
                                bucketOrd = -1 - bucketOrd;
                                collectExistingBucket(sub, doc, bucketOrd);
                            } else {
                                collectBucket(sub, doc, bucketOrd);
                            }
                        }
                        if (endKey > previousKey) {
                            previousKey = endKey;
                        }
                    }
                }
            }
        }
    };
}