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

Example 11 with LeafBucketCollector

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

the class ParentJoinAggregator method beforeBuildingBuckets.

@Override
protected void beforeBuildingBuckets(long[] ordsToCollect) throws IOException {
    IndexReader indexReader = context().searcher().getIndexReader();
    for (LeafReaderContext ctx : indexReader.leaves()) {
        Scorer childDocsScorer = outFilter.scorer(ctx);
        if (childDocsScorer == null) {
            continue;
        }
        DocIdSetIterator childDocsIter = childDocsScorer.iterator();
        final LeafBucketCollector sub = collectableSubAggregators.getLeafCollector(ctx);
        final SortedSetDocValues globalOrdinals = valuesSource.globalOrdinalsValues(ctx);
        // Set the scorer, since we now replay only the child docIds
        sub.setScorer(new Scorable() {

            @Override
            public float score() {
                return 1f;
            }

            @Override
            public int docID() {
                return childDocsIter.docID();
            }
        });
        final Bits liveDocs = ctx.reader().getLiveDocs();
        for (int docId = childDocsIter.nextDoc(); docId != DocIdSetIterator.NO_MORE_DOCS; docId = childDocsIter.nextDoc()) {
            if (liveDocs != null && liveDocs.get(docId) == false) {
                continue;
            }
            if (false == globalOrdinals.advanceExact(docId)) {
                continue;
            }
            int globalOrdinal = (int) globalOrdinals.nextOrd();
            assert globalOrdinal != -1 && globalOrdinals.nextOrd() == SortedSetDocValues.NO_MORE_ORDS;
            /*
                 * Check if we contain every ordinal. It's almost certainly be
                 * faster to replay all the matching ordinals and filter them down
                 * to just those listed in ordsToCollect, but we don't have a data
                 * structure that maps a primitive long to a list of primitive
                 * longs.
                 */
            for (long owningBucketOrd : ordsToCollect) {
                if (collectionStrategy.exists(owningBucketOrd, globalOrdinal)) {
                    collectBucket(sub, docId, owningBucketOrd);
                }
            }
        }
    }
    // Run post collection after collecting the sub-aggs
    super.postCollection();
}

Example 12 with LeafBucketCollector

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

the class CompositeValuesCollectorQueue method getLeafCollector.

/**
 * Creates the collector that will visit the composite buckets of the matching documents.
 * If <code>forceLeadSourceValue</code> is not null, the leading source will use this value
 * for each document.
 * The provided collector <code>in</code> is called on each composite bucket.
 */
LeafBucketCollector getLeafCollector(Comparable forceLeadSourceValue, LeafReaderContext context, LeafBucketCollector in) throws IOException {
    int last = arrays.length - 1;
    LeafBucketCollector collector = in;
    while (last > 0) {
        collector = arrays[last--].getLeafCollector(context, collector);
    }
    if (forceLeadSourceValue != null) {
        collector = arrays[last].getLeafCollector(forceLeadSourceValue, context, collector);
    } else {
        collector = arrays[last].getLeafCollector(context, collector);
    }
    return collector;
}

Example 13 with LeafBucketCollector

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

the class GlobalOrdinalValuesSource method getLeafCollector.

@Override
LeafBucketCollector getLeafCollector(Comparable value, LeafReaderContext context, LeafBucketCollector next) throws IOException {
    if (value.getClass() != BytesRef.class) {
        throw new IllegalArgumentException("Expected BytesRef, got " + value.getClass());
    }
    BytesRef term = (BytesRef) value;
    final SortedSetDocValues dvs = docValuesFunc.apply(context);
    if (lookup == null) {
        initLookup(dvs);
    }
    return new LeafBucketCollector() {

        boolean currentValueIsSet = false;

        @Override
        public void collect(int doc, long bucket) throws IOException {
            if (currentValueIsSet == false) {
                if (dvs.advanceExact(doc)) {
                    long ord;
                    while ((ord = dvs.nextOrd()) != NO_MORE_ORDS) {
                        if (term.equals(lookup.lookupOrd(ord))) {
                            currentValueIsSet = true;
                            currentValue = ord;
                            break;
                        }
                    }
                }
            }
            assert currentValueIsSet;
            next.collect(doc, bucket);
        }
    };
}

Example 14 with LeafBucketCollector

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

the class GlobalOrdinalValuesSource method getLeafCollector.

@Override
LeafBucketCollector getLeafCollector(LeafReaderContext context, LeafBucketCollector next) throws IOException {
    final SortedSetDocValues dvs = docValuesFunc.apply(context);
    if (lookup == null) {
        initLookup(dvs);
    }
    return new LeafBucketCollector() {

        @Override
        public void collect(int doc, long bucket) throws IOException {
            if (dvs.advanceExact(doc)) {
                long ord;
                while ((ord = dvs.nextOrd()) != NO_MORE_ORDS) {
                    currentValue = ord;
                    next.collect(doc, bucket);
                }
            } else if (missingBucket) {
                currentValue = -1;
                next.collect(doc, bucket);
            }
        }
    };
}

Example 15 with LeafBucketCollector

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

the class CompositeAggregator method runDeferredCollections.

/**
 * Replay the documents that might contain a top bucket and pass top buckets to
 * the {@link #deferredCollectors}.
 */
private void runDeferredCollections() throws IOException {
    final boolean needsScores = scoreMode().needsScores();
    Weight weight = null;
    if (needsScores) {
        Query query = context.query();
        weight = context.searcher().createWeight(context.searcher().rewrite(query), ScoreMode.COMPLETE, 1f);
    }
    deferredCollectors.preCollection();
    for (Entry entry : entries) {
        DocIdSetIterator docIdSetIterator = entry.docIdSet.iterator();
        if (docIdSetIterator == null) {
            continue;
        }
        final LeafBucketCollector subCollector = deferredCollectors.getLeafCollector(entry.context);
        final LeafBucketCollector collector = queue.getLeafCollector(entry.context, getSecondPassCollector(subCollector));
        DocIdSetIterator scorerIt = null;
        if (needsScores) {
            Scorer scorer = weight.scorer(entry.context);
            if (scorer != null) {
                scorerIt = scorer.iterator();
                subCollector.setScorer(scorer);
            }
        }
        int docID;
        while ((docID = docIdSetIterator.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {
            if (needsScores) {
                assert scorerIt != null && scorerIt.docID() < docID;
                scorerIt.advance(docID);
                // aggregations should only be replayed on matching documents
                assert scorerIt.docID() == docID;
            }
            collector.collect(docID);
        }
    }
    deferredCollectors.postCollection();
}