Examples with DocIdSetIterator org.apache.lucene.search.DocIdSetIterator used on opensource projects

Example 41 with DocIdSetIterator

use of org.apache.lucene.search.DocIdSetIterator in project elasticsearch by elastic.

the class ProfileScorer method twoPhaseIterator.

@Override
public TwoPhaseIterator twoPhaseIterator() {
    final TwoPhaseIterator in = scorer.twoPhaseIterator();
    if (in == null) {
        return null;
    }
    final DocIdSetIterator inApproximation = in.approximation();
    final DocIdSetIterator approximation = new DocIdSetIterator() {

        @Override
        public int advance(int target) throws IOException {
            profile.startTime(QueryTimingType.ADVANCE);
            try {
                return inApproximation.advance(target);
            } finally {
                profile.stopAndRecordTime();
            }
        }

        @Override
        public int nextDoc() throws IOException {
            profile.startTime(QueryTimingType.NEXT_DOC);
            try {
                return inApproximation.nextDoc();
            } finally {
                profile.stopAndRecordTime();
            }
        }

        @Override
        public int docID() {
            return inApproximation.docID();
        }

        @Override
        public long cost() {
            return inApproximation.cost();
        }
    };
    return new TwoPhaseIterator(approximation) {

        @Override
        public boolean matches() throws IOException {
            profile.startTime(QueryTimingType.MATCH);
            try {
                return in.matches();
            } finally {
                profile.stopAndRecordTime();
            }
        }

        @Override
        public float matchCost() {
            return in.matchCost();
        }
    };
}

Example 42 with DocIdSetIterator

use of org.apache.lucene.search.DocIdSetIterator in project elasticsearch by elastic.

the class TermsSliceQuery method createWeight.

@Override
public Weight createWeight(IndexSearcher searcher, boolean needsScores) throws IOException {
    return new ConstantScoreWeight(this) {

        @Override
        public Scorer scorer(LeafReaderContext context) throws IOException {
            final DocIdSet disi = build(context.reader());
            final DocIdSetIterator leafIt = disi.iterator();
            return new ConstantScoreScorer(this, score(), leafIt);
        }
    };
}

Example 43 with DocIdSetIterator

use of org.apache.lucene.search.DocIdSetIterator in project elasticsearch by elastic.

the class NestedChildrenFilterTests method testNestedChildrenFilter.

public void testNestedChildrenFilter() throws Exception {
    int numParentDocs = scaledRandomIntBetween(0, 32);
    int maxChildDocsPerParent = scaledRandomIntBetween(8, 16);
    Directory dir = newDirectory();
    RandomIndexWriter writer = new RandomIndexWriter(random(), dir);
    for (int i = 0; i < numParentDocs; i++) {
        int numChildDocs = scaledRandomIntBetween(0, maxChildDocsPerParent);
        List<Document> docs = new ArrayList<>(numChildDocs + 1);
        for (int j = 0; j < numChildDocs; j++) {
            Document childDoc = new Document();
            childDoc.add(new StringField("type", "child", Field.Store.NO));
            docs.add(childDoc);
        }
        Document parenDoc = new Document();
        parenDoc.add(new StringField("type", "parent", Field.Store.NO));
        parenDoc.add(new LegacyIntField("num_child_docs", numChildDocs, Field.Store.YES));
        docs.add(parenDoc);
        writer.addDocuments(docs);
    }
    IndexReader reader = writer.getReader();
    writer.close();
    IndexSearcher searcher = new IndexSearcher(reader);
    FetchSubPhase.HitContext hitContext = new FetchSubPhase.HitContext();
    BitSetProducer parentFilter = new QueryBitSetProducer(new TermQuery(new Term("type", "parent")));
    Query childFilter = new TermQuery(new Term("type", "child"));
    int checkedParents = 0;
    final Weight parentsWeight = searcher.createNormalizedWeight(new TermQuery(new Term("type", "parent")), false);
    for (LeafReaderContext leaf : reader.leaves()) {
        DocIdSetIterator parents = parentsWeight.scorer(leaf).iterator();
        for (int parentDoc = parents.nextDoc(); parentDoc != DocIdSetIterator.NO_MORE_DOCS; parentDoc = parents.nextDoc()) {
            int expectedChildDocs = leaf.reader().document(parentDoc).getField("num_child_docs").numericValue().intValue();
            hitContext.reset(null, leaf, parentDoc, searcher);
            NestedChildrenQuery nestedChildrenFilter = new NestedChildrenQuery(parentFilter, childFilter, hitContext);
            TotalHitCountCollector totalHitCountCollector = new TotalHitCountCollector();
            searcher.search(new ConstantScoreQuery(nestedChildrenFilter), totalHitCountCollector);
            assertThat(totalHitCountCollector.getTotalHits(), equalTo(expectedChildDocs));
            checkedParents++;
        }
    }
    assertThat(checkedParents, equalTo(numParentDocs));
    reader.close();
    dir.close();
}

Example 44 with DocIdSetIterator

use of org.apache.lucene.search.DocIdSetIterator in project elasticsearch by elastic.

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.createNormalizedWeight(childFilter, false);
    Scorer childDocsScorer = weight.scorer(ctx);
    final BitSet parentDocs = parentFilter.getBitSet(ctx);
    final DocIdSetIterator childDocs = childDocsScorer != null ? childDocsScorer.iterator() : null;
    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);
            }
        }
    };
}

Example 45 with DocIdSetIterator

use of org.apache.lucene.search.DocIdSetIterator in project elasticsearch by elastic.

the class ParentToChildrenAggregator method doPostCollection.

@Override
protected void doPostCollection() throws IOException {
    IndexReader indexReader = context().searcher().getIndexReader();
    for (LeafReaderContext ctx : indexReader.leaves()) {
        Scorer childDocsScorer = childFilter.scorer(ctx);
        if (childDocsScorer == null) {
            continue;
        }
        DocIdSetIterator childDocsIter = childDocsScorer.iterator();
        final LeafBucketCollector sub = collectableSubAggregators.getLeafCollector(ctx);
        final SortedDocValues globalOrdinals = valuesSource.globalOrdinalsValues(parentType, ctx);
        // Set the scorer, since we now replay only the child docIds
        sub.setScorer(new ConstantScoreScorer(null, 1f, childDocsIter));
        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;
            }
            long globalOrdinal = globalOrdinals.getOrd(docId);
            if (globalOrdinal != -1) {
                long bucketOrd = parentOrdToBuckets.get(globalOrdinal);
                if (bucketOrd != -1) {
                    collectBucket(sub, docId, bucketOrd);
                    if (multipleBucketsPerParentOrd) {
                        long[] otherBucketOrds = parentOrdToOtherBuckets.get(globalOrdinal);
                        if (otherBucketOrds != null) {
                            for (long otherBucketOrd : otherBucketOrds) {
                                collectBucket(sub, docId, otherBucketOrd);
                            }
                        }
                    }
                }
            }
        }
    }
}