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

Example 1 with LeafCollector

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

the class SearchCancellationTests method testCancellableCollector.

public void testCancellableCollector() throws IOException {
    TotalHitCountCollector collector = new TotalHitCountCollector();
    AtomicBoolean cancelled = new AtomicBoolean();
    CancellableCollector cancellableCollector = new CancellableCollector(cancelled::get, false, collector);
    final LeafCollector leafCollector = cancellableCollector.getLeafCollector(reader.leaves().get(0));
    leafCollector.collect(0);
    cancelled.set(true);
    leafCollector.collect(1);
    expectThrows(TaskCancelledException.class, () -> cancellableCollector.getLeafCollector(reader.leaves().get(1)));
}

Example 2 with LeafCollector

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

the class QueryProfilerTests method testCollector.

public void testCollector() throws IOException {
    TotalHitCountCollector collector = new TotalHitCountCollector();
    ProfileCollector profileCollector = new ProfileCollector(collector);
    assertEquals(0, profileCollector.getTime());
    final LeafCollector leafCollector = profileCollector.getLeafCollector(reader.leaves().get(0));
    assertThat(profileCollector.getTime(), greaterThan(0L));
    long time = profileCollector.getTime();
    leafCollector.setScorer(null);
    assertThat(profileCollector.getTime(), greaterThan(time));
    time = profileCollector.getTime();
    leafCollector.collect(0);
    assertThat(profileCollector.getTime(), greaterThan(time));
}

Example 3 with LeafCollector

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

the class ProfileCollector method getLeafCollector.

@Override
public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
    final long start = System.nanoTime();
    final LeafCollector inLeafCollector;
    try {
        inLeafCollector = super.getLeafCollector(context);
    } finally {
        time += Math.max(1, System.nanoTime() - start);
    }
    return new FilterLeafCollector(inLeafCollector) {

        @Override
        public void collect(int doc) throws IOException {
            final long start = System.nanoTime();
            try {
                super.collect(doc);
            } finally {
                time += Math.max(1, System.nanoTime() - start);
            }
        }

        @Override
        public void setScorer(Scorer scorer) throws IOException {
            final long start = System.nanoTime();
            try {
                super.setScorer(scorer);
            } finally {
                time += Math.max(1, System.nanoTime() - start);
            }
        }
    };
}

Example 4 with LeafCollector

use of org.apache.lucene.search.LeafCollector in project lucene-solr by apache.

the class BlockGroupingCollector method getTopGroups.

// TODO: maybe allow no sort on retrieving groups?  app
// may want to simply process docs in the group itself?
// typically they will be presented as a "single" result
// in the UI?
/** Returns the grouped results.  Returns null if the
   *  number of groups collected is <= groupOffset.
   *
   *  <p><b>NOTE</b>: This collector is unable to compute
   *  the groupValue per group so it will always be null.
   *  This is normally not a problem, as you can obtain the
   *  value just like you obtain other values for each
   *  matching document (eg, via stored fields, via
   *  DocValues, etc.)
   *
   *  @param withinGroupSort The {@link Sort} used to sort
   *    documents within each group.
   *  @param groupOffset Which group to start from
   *  @param withinGroupOffset Which document to start from
   *    within each group
   *  @param maxDocsPerGroup How many top documents to keep
   *     within each group.
   *  @param fillSortFields If true then the Comparable
   *     values for the sort fields will be set
   */
public TopGroups<?> getTopGroups(Sort withinGroupSort, int groupOffset, int withinGroupOffset, int maxDocsPerGroup, boolean fillSortFields) throws IOException {
    //}
    if (subDocUpto != 0) {
        processGroup();
    }
    if (groupOffset >= groupQueue.size()) {
        return null;
    }
    int totalGroupedHitCount = 0;
    final FakeScorer fakeScorer = new FakeScorer();
    float maxScore = Float.MIN_VALUE;
    @SuppressWarnings({ "unchecked", "rawtypes" }) final GroupDocs<Object>[] groups = new GroupDocs[groupQueue.size() - groupOffset];
    for (int downTo = groupQueue.size() - groupOffset - 1; downTo >= 0; downTo--) {
        final OneGroup og = groupQueue.pop();
        // At this point we hold all docs w/ in each group,
        // unsorted; we now sort them:
        final TopDocsCollector<?> collector;
        if (withinGroupSort.equals(Sort.RELEVANCE)) {
            // Sort by score
            if (!needsScores) {
                throw new IllegalArgumentException("cannot sort by relevance within group: needsScores=false");
            }
            collector = TopScoreDocCollector.create(maxDocsPerGroup);
        } else {
            // Sort by fields
            collector = TopFieldCollector.create(withinGroupSort, maxDocsPerGroup, fillSortFields, needsScores, needsScores);
        }
        LeafCollector leafCollector = collector.getLeafCollector(og.readerContext);
        leafCollector.setScorer(fakeScorer);
        for (int docIDX = 0; docIDX < og.count; docIDX++) {
            final int doc = og.docs[docIDX];
            fakeScorer.doc = doc;
            if (needsScores) {
                fakeScorer.score = og.scores[docIDX];
            }
            leafCollector.collect(doc);
        }
        totalGroupedHitCount += og.count;
        final Object[] groupSortValues;
        if (fillSortFields) {
            groupSortValues = new Comparable<?>[comparators.length];
            for (int sortFieldIDX = 0; sortFieldIDX < comparators.length; sortFieldIDX++) {
                groupSortValues[sortFieldIDX] = comparators[sortFieldIDX].value(og.comparatorSlot);
            }
        } else {
            groupSortValues = null;
        }
        final TopDocs topDocs = collector.topDocs(withinGroupOffset, maxDocsPerGroup);
        // TODO: we could aggregate scores across children
        // by Sum/Avg instead of passing NaN:
        groups[downTo] = new GroupDocs<>(Float.NaN, topDocs.getMaxScore(), og.count, topDocs.scoreDocs, null, groupSortValues);
        maxScore = Math.max(maxScore, topDocs.getMaxScore());
    }
    return new TopGroups<>(new TopGroups<>(groupSort.getSort(), withinGroupSort.getSort(), totalHitCount, totalGroupedHitCount, groups, maxScore), totalGroupCount);
}

Example 5 with LeafCollector

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

the class FilteredCollector method getLeafCollector.

@Override
public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
    final Scorer filterScorer = filter.scorer(context);
    final LeafCollector in = collector.getLeafCollector(context);
    final Bits bits = Lucene.asSequentialAccessBits(context.reader().maxDoc(), filterScorer);
    return new FilterLeafCollector(in) {

        @Override
        public void collect(int doc) throws IOException {
            if (bits.get(doc)) {
                in.collect(doc);
            }
        }
    };
}