Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Introduce IndexSearcher#searchLeaf(LeafReaderContext, Weight, Collector) method #13603

Merged
merged 3 commits into from
Jul 30, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 45 additions & 28 deletions lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java
Original file line number Diff line number Diff line change
Expand Up @@ -674,7 +674,8 @@ private <C extends Collector, T> T search(
/**
* Lower-level search API.
*
* <p>{@link LeafCollector#collect(int)} is called for every document. <br>
* <p>{@link #searchLeaf(LeafReaderContext, Weight, Collector)} is called for every leaf
* partition. <br>
*
* <p>NOTE: this method executes the searches on all given leaves exclusively. To search across
* all the searchers leaves use {@link #leafContexts}.
Expand All @@ -694,40 +695,56 @@ protected void search(List<LeafReaderContext> leaves, Weight weight, Collector c
// threaded...? the Collector could be sync'd?
// always use single thread:
for (LeafReaderContext ctx : leaves) { // search each subreader
final LeafCollector leafCollector;
searchLeaf(ctx, weight, collector);
}
}

/**
* Lower-level search API
*
* <p>{@link LeafCollector#collect(int)} is called for every document. <br>
*
* @param ctx the leaf to execute the search against
* @param weight to match document
* @param collector to receive hits
* @throws TooManyClauses If a query would exceed {@link IndexSearcher#getMaxClauseCount()}
* clauses.
*/
protected void searchLeaf(LeafReaderContext ctx, Weight weight, Collector collector)
throws IOException {
final LeafCollector leafCollector;
try {
leafCollector = collector.getLeafCollector(ctx);
} catch (
@SuppressWarnings("unused")
CollectionTerminatedException e) {
// there is no doc of interest in this reader context
// continue with the following leaf
return;
}
ScorerSupplier scorerSupplier = weight.scorerSupplier(ctx);
if (scorerSupplier != null) {
scorerSupplier.setTopLevelScoringClause();
BulkScorer scorer = scorerSupplier.bulkScorer();
if (queryTimeout != null) {
scorer = new TimeLimitingBulkScorer(scorer, queryTimeout);
}
try {
leafCollector = collector.getLeafCollector(ctx);
scorer.score(leafCollector, ctx.reader().getLiveDocs());
} catch (
@SuppressWarnings("unused")
CollectionTerminatedException e) {
// there is no doc of interest in this reader context
// collection was terminated prematurely
// continue with the following leaf
continue;
}
ScorerSupplier scorerSupplier = weight.scorerSupplier(ctx);
if (scorerSupplier != null) {
scorerSupplier.setTopLevelScoringClause();
BulkScorer scorer = scorerSupplier.bulkScorer();
if (queryTimeout != null) {
scorer = new TimeLimitingBulkScorer(scorer, queryTimeout);
}
try {
scorer.score(leafCollector, ctx.reader().getLiveDocs());
} catch (
@SuppressWarnings("unused")
CollectionTerminatedException e) {
// collection was terminated prematurely
// continue with the following leaf
} catch (
@SuppressWarnings("unused")
TimeLimitingBulkScorer.TimeExceededException e) {
partialResult = true;
}
} catch (
@SuppressWarnings("unused")
TimeLimitingBulkScorer.TimeExceededException e) {
partialResult = true;
}
// Note: this is called if collection ran successfully, including the above special cases of
// CollectionTerminatedException and TimeExceededException, but no other exception.
leafCollector.finish();
}
// Note: this is called if collection ran successfully, including the above special cases of
// CollectionTerminatedException and TimeExceededException, but no other exception.
leafCollector.finish();
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The diff is misleading: this is a pure cut and paste of the code from within the existing loop, to the new method that's called for every entry.

}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -43,26 +43,26 @@
public class TestTopDocsMerge extends LuceneTestCase {

private static class ShardSearcher extends IndexSearcher {
private final List<LeafReaderContext> ctx;
private final LeafReaderContext ctx;

public ShardSearcher(LeafReaderContext ctx, IndexReaderContext parent) {
super(parent);
this.ctx = Collections.singletonList(ctx);
this.ctx = ctx;
}

public void search(Weight weight, Collector collector) throws IOException {
search(ctx, weight, collector);
searchLeaf(ctx, weight, collector);
}

public TopDocs search(Weight weight, int topN) throws IOException {
TopScoreDocCollector collector = TopScoreDocCollector.create(topN, Integer.MAX_VALUE);
search(ctx, weight, collector);
searchLeaf(ctx, weight, collector);
return collector.topDocs();
}

@Override
public String toString() {
return "ShardSearcher(" + ctx.get(0) + ")";
return "ShardSearcher(" + ctx + ")";
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1571,20 +1571,20 @@ private void assertEquals(
}

private static class ShardSearcher extends IndexSearcher {
private final List<LeafReaderContext> ctx;
private final LeafReaderContext ctx;

public ShardSearcher(LeafReaderContext ctx, IndexReaderContext parent) {
super(parent);
this.ctx = Collections.singletonList(ctx);
this.ctx = ctx;
}

public void search(Weight weight, Collector collector) throws IOException {
search(ctx, weight, collector);
searchLeaf(ctx, weight, collector);
}

@Override
public String toString() {
return "ShardSearcher(" + ctx.get(0).reader() + ")";
return "ShardSearcher(" + ctx.reader() + ")";
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
package org.apache.lucene.tests.search;

import java.io.IOException;
import java.util.List;
import java.util.concurrent.Executor;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.LeafReaderContext;
Expand Down Expand Up @@ -53,25 +52,22 @@ public ScorerIndexSearcher(IndexReader r) {
}

@Override
protected void search(List<LeafReaderContext> leaves, Weight weight, Collector collector)
protected void searchLeaf(LeafReaderContext ctx, Weight weight, Collector collector)
throws IOException {
collector.setWeight(weight);
for (LeafReaderContext ctx : leaves) { // search each subreader
// we force the use of Scorer (not BulkScorer) to make sure
// that the scorer passed to LeafCollector.setScorer supports
// Scorer.getChildren
Scorer scorer = weight.scorer(ctx);
if (scorer != null) {
final DocIdSetIterator iterator = scorer.iterator();
final LeafCollector leafCollector = collector.getLeafCollector(ctx);
leafCollector.setScorer(scorer);
final Bits liveDocs = ctx.reader().getLiveDocs();
for (int doc = iterator.nextDoc();
doc != DocIdSetIterator.NO_MORE_DOCS;
doc = iterator.nextDoc()) {
if (liveDocs == null || liveDocs.get(doc)) {
leafCollector.collect(doc);
}
// we force the use of Scorer (not BulkScorer) to make sure
// that the scorer passed to LeafCollector.setScorer supports
// Scorer.getChildren
Scorer scorer = weight.scorer(ctx);
if (scorer != null) {
final DocIdSetIterator iterator = scorer.iterator();
final LeafCollector leafCollector = collector.getLeafCollector(ctx);
leafCollector.setScorer(scorer);
final Bits liveDocs = ctx.reader().getLiveDocs();
for (int doc = iterator.nextDoc();
doc != DocIdSetIterator.NO_MORE_DOCS;
doc = iterator.nextDoc()) {
if (liveDocs == null || liveDocs.get(doc)) {
leafCollector.collect(doc);
}
}
}
Expand Down