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

Example 1 with LeafBucketCollectorBase

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

the class MatrixStatsAggregator 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 NumericDoubleValues[] values = new NumericDoubleValues[valuesSources.fieldNames().length];
    for (int i = 0; i < values.length; ++i) {
        values[i] = valuesSources.getField(i, ctx);
    }
    return new LeafBucketCollectorBase(sub, values) {

        final String[] fieldNames = valuesSources.fieldNames();

        final double[] fieldVals = new double[fieldNames.length];

        @Override
        public void collect(int doc, long bucket) throws IOException {
            // get fields
            if (includeDocument(doc)) {
                stats = bigArrays.grow(stats, bucket + 1);
                RunningStats stat = stats.get(bucket);
                // add document fields to correlation stats
                if (stat == null) {
                    stat = new RunningStats(fieldNames, fieldVals);
                    stats.set(bucket, stat);
                } else {
                    stat.add(fieldNames, fieldVals);
                }
            }
        }

        /**
         * return a map of field names and data
         */
        private boolean includeDocument(int doc) throws IOException {
            // loop over fields
            for (int i = 0; i < fieldVals.length; ++i) {
                final NumericDoubleValues doubleValues = values[i];
                if (doubleValues.advanceExact(doc)) {
                    final double value = doubleValues.doubleValue();
                    if (value == Double.NEGATIVE_INFINITY) {
                        // TODO: Fix matrix stats to treat neg inf as any other value
                        return false;
                    }
                    fieldVals[i] = value;
                } else {
                    return false;
                }
            }
            return true;
        }
    };
}

Example 2 with LeafBucketCollectorBase

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

the class GeoGridAggregator method getLeafCollector.

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

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

Example 3 with LeafBucketCollectorBase

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

the class GlobalOrdinalsStringTermsAggregator method getLeafCollector.

@Override
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
    SortedSetDocValues globalOrds = valuesSource.globalOrdinalsValues(ctx);
    collectionStrategy.globalOrdsReady(globalOrds);
    SortedDocValues singleValues = DocValues.unwrapSingleton(globalOrds);
    if (singleValues != null) {
        segmentsWithSingleValuedOrds++;
        if (acceptedGlobalOrdinals == ALWAYS_TRUE) {
            /*
                 * Optimize when there isn't a filter because that is very
                 * common and marginally faster.
                 */
            return resultStrategy.wrapCollector(new LeafBucketCollectorBase(sub, globalOrds) {

                @Override
                public void collect(int doc, long owningBucketOrd) throws IOException {
                    if (false == singleValues.advanceExact(doc)) {
                        return;
                    }
                    int globalOrd = singleValues.ordValue();
                    collectionStrategy.collectGlobalOrd(owningBucketOrd, doc, globalOrd, sub);
                }
            });
        }
        return resultStrategy.wrapCollector(new LeafBucketCollectorBase(sub, globalOrds) {

            @Override
            public void collect(int doc, long owningBucketOrd) throws IOException {
                if (false == singleValues.advanceExact(doc)) {
                    return;
                }
                int globalOrd = singleValues.ordValue();
                if (false == acceptedGlobalOrdinals.test(globalOrd)) {
                    return;
                }
                collectionStrategy.collectGlobalOrd(owningBucketOrd, doc, globalOrd, sub);
            }
        });
    }
    segmentsWithMultiValuedOrds++;
    if (acceptedGlobalOrdinals == ALWAYS_TRUE) {
        /*
             * Optimize when there isn't a filter because that is very
             * common and marginally faster.
             */
        return resultStrategy.wrapCollector(new LeafBucketCollectorBase(sub, globalOrds) {

            @Override
            public void collect(int doc, long owningBucketOrd) throws IOException {
                if (false == globalOrds.advanceExact(doc)) {
                    return;
                }
                for (long globalOrd = globalOrds.nextOrd(); globalOrd != NO_MORE_ORDS; globalOrd = globalOrds.nextOrd()) {
                    collectionStrategy.collectGlobalOrd(owningBucketOrd, doc, globalOrd, sub);
                }
            }
        });
    }
    return resultStrategy.wrapCollector(new LeafBucketCollectorBase(sub, globalOrds) {

        @Override
        public void collect(int doc, long owningBucketOrd) throws IOException {
            if (false == globalOrds.advanceExact(doc)) {
                return;
            }
            for (long globalOrd = globalOrds.nextOrd(); globalOrd != NO_MORE_ORDS; globalOrd = globalOrds.nextOrd()) {
                if (false == acceptedGlobalOrdinals.test(globalOrd)) {
                    continue;
                }
                collectionStrategy.collectGlobalOrd(owningBucketOrd, doc, globalOrd, sub);
            }
        }
    });
}

Example 4 with LeafBucketCollectorBase

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

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 {
            if (values.advanceExact(doc)) {
                final int valuesCount = values.docValueCount();
                for (int i = 0, lo = 0; i < valuesCount; ++i) {
                    final double value = values.nextValue();
                    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 5 with LeafBucketCollectorBase

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

the class SumAggregator 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 {
            sums = bigArrays.grow(sums, bucket + 1);
            compensations = bigArrays.grow(compensations, bucket + 1);
            if (values.advanceExact(doc)) {
                final int valuesCount = values.docValueCount();
                // 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 < valuesCount; i++) {
                    double value = values.nextValue();
                    kahanSummation.add(value);
                }
                compensations.set(bucket, kahanSummation.delta());
                sums.set(bucket, kahanSummation.value());
            }
        }
    };
}