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

Example 21 with LeafBucketCollectorBase

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

the class AbstractTDigestPercentilesAggregator 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 {
            TDigestState state = getExistingOrNewHistogram(bigArrays, bucket);
            if (values.advanceExact(doc)) {
                final int valueCount = values.docValueCount();
                for (int i = 0; i < valueCount; i++) {
                    state.add(values.nextValue());
                }
            }
        }
    };
}

Example 22 with LeafBucketCollectorBase

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

the class LongRareTermsAggregator method getLeafCollector.

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

        @Override
        public void collect(int docId, long owningBucketOrd) throws IOException {
            if (false == values.advanceExact(docId)) {
                return;
            }
            int valuesCount = values.docValueCount();
            long previous = Long.MAX_VALUE;
            for (int i = 0; i < valuesCount; ++i) {
                long val = values.nextValue();
                if (i == 0 && previous == val) {
                    continue;
                }
                previous = val;
                if (filter != null && false == filter.accept(val)) {
                    continue;
                }
                long bucketOrdinal = bucketOrds.add(owningBucketOrd, val);
                if (bucketOrdinal < 0) {
                    // already seen
                    bucketOrdinal = -1 - bucketOrdinal;
                    collectExistingBucket(sub, docId, bucketOrdinal);
                } else {
                    collectBucket(sub, docId, bucketOrdinal);
                }
            }
        }
    };
}

Example 23 with LeafBucketCollectorBase

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

the class ValueCountAggregator 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();
    if (valuesSource instanceof ValuesSource.Numeric) {
        final SortedNumericDocValues values = ((ValuesSource.Numeric) valuesSource).longValues(ctx);
        return new LeafBucketCollectorBase(sub, values) {

            @Override
            public void collect(int doc, long bucket) throws IOException {
                counts = bigArrays.grow(counts, bucket + 1);
                if (values.advanceExact(doc)) {
                    counts.increment(bucket, values.docValueCount());
                }
            }
        };
    }
    if (valuesSource instanceof ValuesSource.Bytes.GeoPoint) {
        MultiGeoPointValues values = ((ValuesSource.GeoPoint) valuesSource).geoPointValues(ctx);
        return new LeafBucketCollectorBase(sub, null) {

            @Override
            public void collect(int doc, long bucket) throws IOException {
                counts = bigArrays.grow(counts, bucket + 1);
                if (values.advanceExact(doc)) {
                    counts.increment(bucket, values.docValueCount());
                }
            }
        };
    }
    // The following is default collector. Including the keyword FieldType
    final SortedBinaryDocValues values = valuesSource.bytesValues(ctx);
    return new LeafBucketCollectorBase(sub, values) {

        @Override
        public void collect(int doc, long bucket) throws IOException {
            counts = bigArrays.grow(counts, bucket + 1);
            if (values.advanceExact(doc)) {
                counts.increment(bucket, values.docValueCount());
            }
        }
    };
}

Example 24 with LeafBucketCollectorBase

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

the class WeightedAvgAggregator method getLeafCollector.

@Override
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
    if (valuesSources == null) {
        return LeafBucketCollector.NO_OP_COLLECTOR;
    }
    final BigArrays bigArrays = context.bigArrays();
    final SortedNumericDoubleValues docValues = valuesSources.getField(VALUE_FIELD.getPreferredName(), ctx);
    final SortedNumericDoubleValues docWeights = valuesSources.getField(WEIGHT_FIELD.getPreferredName(), ctx);
    final CompensatedSum compensatedValueSum = new CompensatedSum(0, 0);
    final CompensatedSum compensatedWeightSum = new CompensatedSum(0, 0);
    return new LeafBucketCollectorBase(sub, docValues) {

        @Override
        public void collect(int doc, long bucket) throws IOException {
            weights = bigArrays.grow(weights, bucket + 1);
            valueSums = bigArrays.grow(valueSums, bucket + 1);
            valueCompensations = bigArrays.grow(valueCompensations, bucket + 1);
            weightCompensations = bigArrays.grow(weightCompensations, bucket + 1);
            if (docValues.advanceExact(doc) && docWeights.advanceExact(doc)) {
                if (docWeights.docValueCount() > 1) {
                    throw new AggregationExecutionException("Encountered more than one weight for a " + "single document. Use a script to combine multiple weights-per-doc into a single value.");
                }
                // a real weight or a `missing` weight
                assert docWeights.docValueCount() == 1;
                final double weight = docWeights.nextValue();
                final int numValues = docValues.docValueCount();
                assert numValues > 0;
                double valueSum = valueSums.get(bucket);
                double valueCompensation = valueCompensations.get(bucket);
                compensatedValueSum.reset(valueSum, valueCompensation);
                double weightSum = weights.get(bucket);
                double weightCompensation = weightCompensations.get(bucket);
                compensatedWeightSum.reset(weightSum, weightCompensation);
                for (int i = 0; i < numValues; i++) {
                    compensatedValueSum.add(docValues.nextValue() * weight);
                    compensatedWeightSum.add(weight);
                }
                valueSums.set(bucket, compensatedValueSum.value());
                valueCompensations.set(bucket, compensatedValueSum.delta());
                weights.set(bucket, compensatedWeightSum.value());
                weightCompensations.set(bucket, compensatedWeightSum.delta());
            }
        }
    };
}

Example 25 with LeafBucketCollectorBase

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

the class MaxAggregator method getLeafCollector.

@Override
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
    if (valuesSource == null) {
        if (parent != null) {
            return LeafBucketCollector.NO_OP_COLLECTOR;
        } else {
            // we have no parent and the values source is empty so we can skip collecting hits.
            throw new CollectionTerminatedException();
        }
    }
    if (pointConverter != null) {
        Number segMax = findLeafMaxValue(ctx.reader(), pointField, pointConverter);
        if (segMax != null) {
            /*
                 * There is no parent aggregator (see {@link AggregatorBase#getPointReaderOrNull}
                 * so the ordinal for the bucket is always 0.
                 */
            assert maxes.size() == 1;
            double max = maxes.get(0);
            max = Math.max(max, segMax.doubleValue());
            maxes.set(0, max);
            // the maximum value has been extracted, we don't need to collect hits on this segment.
            throw new CollectionTerminatedException();
        }
    }
    final BigArrays bigArrays = context.bigArrays();
    final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx);
    final NumericDoubleValues values = MultiValueMode.MAX.select(allValues);
    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);
            }
            if (values.advanceExact(doc)) {
                final double value = values.doubleValue();
                double max = maxes.get(bucket);
                max = Math.max(max, value);
                maxes.set(bucket, max);
            }
        }
    };
}