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

There's a couple of places in the codebase where we extend `IndexSearcher` to customize
per leaf behaviour, and in order to do that, we need to override the entire search method
that loops through the leaves. A good example is `ScorerIndexSearcher`.

Adding a `searchLeaf` method that provides the per leaf behaviour makes those cases a little
easier to deal with.
This commit is contained in:
Luca Cavanna 2024-07-30 17:17:27 +02:00 committed by GitHub
parent 68aa629f6c
commit 30c965ea57
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 73 additions and 56 deletions

View File

@ -259,6 +259,10 @@ API Changes
* GITHUB#13612: Hunspell: add Suggester#proceedPastRep to avoid losing relevant suggestions. (Peter Gromov)
* GITHUB#13603: Introduced `IndexSearcher#searchLeaf(LeafReaderContext, Weight, Collector)` protected method to
facilitate customizing per-leaf behavior of search without requiring to override
`search(LeafReaderContext[], Weight, Collector)` which requires overriding the entire loop across the leaves (Luca Cavanna)
New Features
---------------------

View File

@ -674,7 +674,8 @@ public class IndexSearcher {
/**
* 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}.
@ -694,40 +695,56 @@ public class IndexSearcher {
// 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;
} catch (
@SuppressWarnings("unused")
TimeLimitingBulkScorer.TimeExceededException e) {
partialResult = true;
}
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;
}
}
// 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();
}
/**

View File

@ -43,27 +43,27 @@ import org.apache.lucene.util.BytesRef;
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 =
new TopScoreDocCollectorManager(topN, null, Integer.MAX_VALUE, false).newCollector();
search(ctx, weight, collector);
searchLeaf(ctx, weight, collector);
return collector.topDocs();
}
@Override
public String toString() {
return "ShardSearcher(" + ctx.get(0) + ")";
return "ShardSearcher(" + ctx + ")";
}
}

View File

@ -1571,20 +1571,20 @@ public class TestGrouping 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);
}
@Override
public String toString() {
return "ShardSearcher(" + ctx.get(0).reader() + ")";
return "ShardSearcher(" + ctx.reader() + ")";
}
}
}

View File

@ -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;
@ -53,25 +52,22 @@ public class ScorerIndexSearcher extends IndexSearcher {
}
@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);
}
}
}