From de07bdd88dd124da1da8709e337f2b5839f6207d Mon Sep 17 00:00:00 2001 From: Adrien Grand Date: Fri, 30 Jun 2023 15:19:35 +0200 Subject: [PATCH] Add a post-collection hook to LeafCollector. (#12380) This adds `LeafCollector#finish` as a per-segment post-collection hook. While it was already possible to do this sort of things on top of the collector API before, a downside is that the last leaf would need to be post-collected in the current thread instead of using the executor, which is a missed opportunity for making queries concurrent. --- lucene/CHANGES.txt | 4 +- .../lucene/search/CachingCollector.java | 71 +++++++++---------- .../lucene/search/FilterLeafCollector.java | 5 ++ .../apache/lucene/search/IndexSearcher.java | 3 + .../apache/lucene/search/LeafCollector.java | 9 +++ .../apache/lucene/search/MultiCollector.java | 10 +++ .../lucene/search/TestCachingCollector.java | 2 + .../lucene/facet/DrillSidewaysScorer.java | 17 +++++ .../apache/lucene/facet/FacetsCollector.java | 19 +++-- .../grouping/BlockGroupingCollector.java | 13 ++-- .../search/grouping/GroupFacetCollector.java | 11 +-- .../grouping/TermGroupFacetCollector.java | 8 --- .../document/SuggestIndexSearcher.java | 5 +- .../document/TopSuggestDocsCollector.java | 19 +++-- .../tests/search/AssertingCollector.java | 13 +++- .../tests/search/AssertingIndexSearcher.java | 4 +- .../tests/search/AssertingLeafCollector.java | 8 +++ 17 files changed, 141 insertions(+), 80 deletions(-) diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt index ef5c9e496a19..b931b32eda24 100644 --- a/lucene/CHANGES.txt +++ b/lucene/CHANGES.txt @@ -11,7 +11,9 @@ API Changes New Features --------------------- -(No changes) + +* GITHUB#12383: Introduced LeafCollector#finish, a hook that runs after + collection has finished running on a leaf. (Adrien Grand) Improvements --------------------- diff --git a/lucene/core/src/java/org/apache/lucene/search/CachingCollector.java b/lucene/core/src/java/org/apache/lucene/search/CachingCollector.java index d065efd3406f..fd45666d33ae 100644 --- a/lucene/core/src/java/org/apache/lucene/search/CachingCollector.java +++ b/lucene/core/src/java/org/apache/lucene/search/CachingCollector.java @@ -66,7 +66,6 @@ private static class NoScoreCachingCollector extends CachingCollector { List contexts; List docs; int maxDocsToCache; - NoScoreCachingLeafCollector lastCollector; NoScoreCachingCollector(Collector in, int maxDocsToCache) { super(in); @@ -76,7 +75,7 @@ private static class NoScoreCachingCollector extends CachingCollector { } protected NoScoreCachingLeafCollector wrap(LeafCollector in, int maxDocsToCache) { - return new NoScoreCachingLeafCollector(in, maxDocsToCache); + return new NoScoreCachingLeafCollector(in, maxDocsToCache, this); } // note: do *not* override needScore to say false. Just because we aren't caching the score @@ -85,13 +84,12 @@ protected NoScoreCachingLeafCollector wrap(LeafCollector in, int maxDocsToCache) @Override public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException { - postCollection(); final LeafCollector in = this.in.getLeafCollector(context); - if (contexts != null) { - contexts.add(context); - } if (maxDocsToCache >= 0) { - return lastCollector = wrap(in, maxDocsToCache); + if (contexts != null) { + contexts.add(context); + } + return wrap(in, maxDocsToCache); } else { return in; } @@ -103,33 +101,16 @@ protected void invalidate() { this.docs = null; } - protected void postCollect(NoScoreCachingLeafCollector collector) { - final int[] docs = collector.cachedDocs(); - maxDocsToCache -= docs.length; - this.docs.add(docs); - } - - private void postCollection() { - if (lastCollector != null) { - if (!lastCollector.hasCache()) { - invalidate(); - } else { - postCollect(lastCollector); - } - lastCollector = null; - } - } - protected void collect(LeafCollector collector, int i) throws IOException { final int[] docs = this.docs.get(i); for (int doc : docs) { collector.collect(doc); } + collector.finish(); } @Override public void replay(Collector other) throws IOException { - postCollection(); if (!isCached()) { throw new IllegalStateException( "cannot replay: cache was cleared because too much RAM was required"); @@ -154,14 +135,7 @@ private static class ScoreCachingCollector extends NoScoreCachingCollector { @Override protected NoScoreCachingLeafCollector wrap(LeafCollector in, int maxDocsToCache) { - return new ScoreCachingLeafCollector(in, maxDocsToCache); - } - - @Override - protected void postCollect(NoScoreCachingLeafCollector collector) { - final ScoreCachingLeafCollector coll = (ScoreCachingLeafCollector) collector; - super.postCollect(coll); - scores.add(coll.cachedScores()); + return new ScoreCachingLeafCollector(in, maxDocsToCache, this); } /** @@ -191,12 +165,15 @@ protected void collect(LeafCollector collector, int i) throws IOException { private class NoScoreCachingLeafCollector extends FilterLeafCollector { final int maxDocsToCache; + final NoScoreCachingCollector collector; int[] docs; int docCount; - NoScoreCachingLeafCollector(LeafCollector in, int maxDocsToCache) { + NoScoreCachingLeafCollector( + LeafCollector in, int maxDocsToCache, NoScoreCachingCollector collector) { super(in); this.maxDocsToCache = maxDocsToCache; + this.collector = collector; docs = new int[Math.min(maxDocsToCache, INITIAL_ARRAY_SIZE)]; docCount = 0; } @@ -235,6 +212,21 @@ public void collect(int doc) throws IOException { super.collect(doc); } + protected void postCollect() { + final int[] docs = cachedDocs(); + collector.maxDocsToCache -= docs.length; + collector.docs.add(docs); + } + + @Override + public void finish() { + if (!hasCache()) { + collector.invalidate(); + } else { + postCollect(); + } + } + boolean hasCache() { return docs != null; } @@ -249,8 +241,9 @@ private class ScoreCachingLeafCollector extends NoScoreCachingLeafCollector { Scorable scorer; float[] scores; - ScoreCachingLeafCollector(LeafCollector in, int maxDocsToCache) { - super(in, maxDocsToCache); + ScoreCachingLeafCollector( + LeafCollector in, int maxDocsToCache, ScoreCachingCollector collector) { + super(in, maxDocsToCache, collector); scores = new float[docs.length]; } @@ -281,6 +274,12 @@ protected void buffer(int doc) throws IOException { float[] cachedScores() { return docs == null ? null : ArrayUtil.copyOfSubArray(scores, 0, docCount); } + + @Override + protected void postCollect() { + super.postCollect(); + ((ScoreCachingCollector) collector).scores.add(cachedScores()); + } } /** diff --git a/lucene/core/src/java/org/apache/lucene/search/FilterLeafCollector.java b/lucene/core/src/java/org/apache/lucene/search/FilterLeafCollector.java index 24733668ff6c..d9bc671fd42a 100644 --- a/lucene/core/src/java/org/apache/lucene/search/FilterLeafCollector.java +++ b/lucene/core/src/java/org/apache/lucene/search/FilterLeafCollector.java @@ -42,6 +42,11 @@ public void collect(int doc) throws IOException { in.collect(doc); } + @Override + public void finish() throws IOException { + in.finish(); + } + @Override public String toString() { String name = getClass().getSimpleName(); diff --git a/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java b/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java index 9c0a29c052fe..b76bf3161d07 100644 --- a/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java +++ b/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java @@ -782,6 +782,9 @@ protected void search(List leaves, Weight weight, Collector c 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(); } } diff --git a/lucene/core/src/java/org/apache/lucene/search/LeafCollector.java b/lucene/core/src/java/org/apache/lucene/search/LeafCollector.java index a42d531c2b29..334afc798ca6 100644 --- a/lucene/core/src/java/org/apache/lucene/search/LeafCollector.java +++ b/lucene/core/src/java/org/apache/lucene/search/LeafCollector.java @@ -95,4 +95,13 @@ public interface LeafCollector { default DocIdSetIterator competitiveIterator() throws IOException { return null; } + + /** + * Hook that gets called once the leaf that is associated with this collector has finished + * collecting successfully, including when a {@link CollectionTerminatedException} is thrown. This + * is typically useful to compile data that has been collected on this leaf, e.g. to convert facet + * counts on leaf ordinals to facet counts on global ordinals. The default implementation does + * nothing. + */ + default void finish() throws IOException {} } diff --git a/lucene/core/src/java/org/apache/lucene/search/MultiCollector.java b/lucene/core/src/java/org/apache/lucene/search/MultiCollector.java index 7f8e81217349..ff6a6a97ba5b 100644 --- a/lucene/core/src/java/org/apache/lucene/search/MultiCollector.java +++ b/lucene/core/src/java/org/apache/lucene/search/MultiCollector.java @@ -223,6 +223,7 @@ public void collect(int doc) throws IOException { } catch ( @SuppressWarnings("unused") CollectionTerminatedException e) { + collectors[i].finish(); collectors[i] = null; if (allCollectorsTerminated()) { throw new CollectionTerminatedException(); @@ -232,6 +233,15 @@ public void collect(int doc) throws IOException { } } + @Override + public void finish() throws IOException { + for (LeafCollector collector : collectors) { + if (collector != null) { + collector.finish(); + } + } + } + private boolean allCollectorsTerminated() { for (int i = 0; i < collectors.length; i++) { if (collectors[i] != null) { diff --git a/lucene/core/src/test/org/apache/lucene/search/TestCachingCollector.java b/lucene/core/src/test/org/apache/lucene/search/TestCachingCollector.java index 4b07e06dd0a4..425c33a34e7c 100644 --- a/lucene/core/src/test/org/apache/lucene/search/TestCachingCollector.java +++ b/lucene/core/src/test/org/apache/lucene/search/TestCachingCollector.java @@ -57,6 +57,7 @@ public void testBasic() throws Exception { for (int i = 0; i < 1000; i++) { acc.collect(i); } + acc.finish(); // now replay them cc.replay( @@ -127,6 +128,7 @@ public void testNoWrappedCollector() throws Exception { acc.collect(0); assertTrue(cc.isCached()); + acc.finish(); cc.replay(new NoOpCollector()); } } diff --git a/lucene/facet/src/java/org/apache/lucene/facet/DrillSidewaysScorer.java b/lucene/facet/src/java/org/apache/lucene/facet/DrillSidewaysScorer.java index c86432b30f7f..ad82594f41db 100644 --- a/lucene/facet/src/java/org/apache/lucene/facet/DrillSidewaysScorer.java +++ b/lucene/facet/src/java/org/apache/lucene/facet/DrillSidewaysScorer.java @@ -18,6 +18,7 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collection; import java.util.Collections; import java.util.Comparator; @@ -198,6 +199,7 @@ private void doQueryFirstScoringSingleDim( docID = baseApproximation.nextDoc(); } + finish(collector, Collections.singleton(dim)); } /** @@ -334,6 +336,8 @@ protected boolean lessThan(DocsAndCost a, DocsAndCost b) { docID = baseApproximation.nextDoc(); } + + finish(collector, sidewaysDims); } private static int advanceIfBehind(int docID, DocIdSetIterator iterator) throws IOException { @@ -552,6 +556,7 @@ private void doDrillDownAdvanceScoring( nextChunkStart += CHUNK; } + finish(collector, Arrays.asList(dims)); } private void doUnionScoring(Bits acceptDocs, LeafCollector collector, DocsAndCost[] dims) @@ -706,6 +711,8 @@ private void doUnionScoring(Bits acceptDocs, LeafCollector collector, DocsAndCos nextChunkStart += CHUNK; } + + finish(collector, Arrays.asList(dims)); } private void collectHit(LeafCollector collector, DocsAndCost[] dims) throws IOException { @@ -757,6 +764,16 @@ private void collectNearMiss(LeafCollector sidewaysCollector) throws IOException sidewaysCollector.collect(collectDocID); } + private void finish(LeafCollector collector, Collection dims) throws IOException { + collector.finish(); + if (drillDownLeafCollector != null) { + drillDownLeafCollector.finish(); + } + for (DocsAndCost dim : dims) { + dim.sidewaysLeafCollector.finish(); + } + } + private void setScorer(LeafCollector mainCollector, Scorable scorer) throws IOException { mainCollector.setScorer(scorer); if (drillDownLeafCollector != null) { diff --git a/lucene/facet/src/java/org/apache/lucene/facet/FacetsCollector.java b/lucene/facet/src/java/org/apache/lucene/facet/FacetsCollector.java index cf49aef39eef..2bce78e22b3c 100644 --- a/lucene/facet/src/java/org/apache/lucene/facet/FacetsCollector.java +++ b/lucene/facet/src/java/org/apache/lucene/facet/FacetsCollector.java @@ -103,13 +103,6 @@ public final boolean getKeepScores() { /** Returns the documents matched by the query, one {@link MatchingDocs} per visited segment. */ public List getMatchingDocs() { - if (docsBuilder != null) { - matchingDocs.add(new MatchingDocs(this.context, docsBuilder.build(), totalHits, scores)); - docsBuilder = null; - scores = null; - context = null; - } - return matchingDocs; } @@ -139,9 +132,7 @@ public final void setScorer(Scorable scorer) throws IOException { @Override protected void doSetNextReader(LeafReaderContext context) throws IOException { - if (docsBuilder != null) { - matchingDocs.add(new MatchingDocs(this.context, docsBuilder.build(), totalHits, scores)); - } + assert docsBuilder == null; docsBuilder = new DocIdSetBuilder(context.reader().maxDoc()); totalHits = 0; if (keepScores) { @@ -150,6 +141,14 @@ protected void doSetNextReader(LeafReaderContext context) throws IOException { this.context = context; } + @Override + public void finish() throws IOException { + matchingDocs.add(new MatchingDocs(this.context, docsBuilder.build(), totalHits, scores)); + docsBuilder = null; + scores = null; + context = null; + } + /** Utility method, to search and also collect all hits into the provided {@link Collector}. */ public static TopDocs search(IndexSearcher searcher, Query q, int n, Collector fc) throws IOException { diff --git a/lucene/grouping/src/java/org/apache/lucene/search/grouping/BlockGroupingCollector.java b/lucene/grouping/src/java/org/apache/lucene/search/grouping/BlockGroupingCollector.java index 9ead686831e9..26c3c915dd37 100644 --- a/lucene/grouping/src/java/org/apache/lucene/search/grouping/BlockGroupingCollector.java +++ b/lucene/grouping/src/java/org/apache/lucene/search/grouping/BlockGroupingCollector.java @@ -270,9 +270,6 @@ public TopGroups getTopGroups( // if (queueFull) { // System.out.println("getTopGroups groupOffset=" + groupOffset + " topNGroups=" + topNGroups); // } - if (subDocUpto != 0) { - processGroup(); - } if (groupOffset >= groupQueue.size()) { return null; } @@ -472,9 +469,6 @@ public void collect(int doc) throws IOException { @Override protected void doSetNextReader(LeafReaderContext readerContext) throws IOException { - if (subDocUpto != 0) { - processGroup(); - } subDocUpto = 0; docBase = readerContext.docBase; // System.out.println("setNextReader base=" + docBase + " r=" + readerContext.reader); @@ -492,6 +486,13 @@ protected void doSetNextReader(LeafReaderContext readerContext) throws IOExcepti } } + @Override + public void finish() throws IOException { + if (subDocUpto != 0) { + processGroup(); + } + } + @Override public ScoreMode scoreMode() { return needsScores ? ScoreMode.COMPLETE : ScoreMode.COMPLETE_NO_SCORES; diff --git a/lucene/grouping/src/java/org/apache/lucene/search/grouping/GroupFacetCollector.java b/lucene/grouping/src/java/org/apache/lucene/search/grouping/GroupFacetCollector.java index 74a957d809dd..4e56a12c9024 100644 --- a/lucene/grouping/src/java/org/apache/lucene/search/grouping/GroupFacetCollector.java +++ b/lucene/grouping/src/java/org/apache/lucene/search/grouping/GroupFacetCollector.java @@ -67,11 +67,6 @@ protected GroupFacetCollector(String groupField, String facetField, BytesRef fac */ public GroupedFacetResult mergeSegmentResults(int size, int minCount, boolean orderByCount) throws IOException { - if (segmentFacetCounts != null) { - segmentResults.add(createSegmentResult()); - segmentFacetCounts = null; // reset - } - int totalCount = 0; int missingCount = 0; SegmentResultPriorityQueue segments = new SegmentResultPriorityQueue(segmentResults.size()); @@ -109,6 +104,12 @@ public GroupedFacetResult mergeSegmentResults(int size, int minCount, boolean or return facetResult; } + @Override + public void finish() throws IOException { + segmentResults.add(createSegmentResult()); + segmentFacetCounts = null; + } + protected abstract SegmentResult createSegmentResult() throws IOException; @Override diff --git a/lucene/grouping/src/java/org/apache/lucene/search/grouping/TermGroupFacetCollector.java b/lucene/grouping/src/java/org/apache/lucene/search/grouping/TermGroupFacetCollector.java index c1b49758b133..e49e517faa88 100644 --- a/lucene/grouping/src/java/org/apache/lucene/search/grouping/TermGroupFacetCollector.java +++ b/lucene/grouping/src/java/org/apache/lucene/search/grouping/TermGroupFacetCollector.java @@ -141,10 +141,6 @@ public void collect(int doc) throws IOException { @Override protected void doSetNextReader(LeafReaderContext context) throws IOException { - if (segmentFacetCounts != null) { - segmentResults.add(createSegmentResult()); - } - groupFieldTermsIndex = DocValues.getSorted(context.reader(), groupField); facetFieldTermsIndex = DocValues.getSorted(context.reader(), facetField); @@ -321,10 +317,6 @@ private void process(int groupOrd, int facetOrd) throws IOException { @Override protected void doSetNextReader(LeafReaderContext context) throws IOException { - if (segmentFacetCounts != null) { - segmentResults.add(createSegmentResult()); - } - groupFieldTermsIndex = DocValues.getSorted(context.reader(), groupField); facetFieldDocTermOrds = DocValues.getSortedSet(context.reader(), facetField); facetFieldNumTerms = (int) facetFieldDocTermOrds.getValueCount(); diff --git a/lucene/suggest/src/java/org/apache/lucene/search/suggest/document/SuggestIndexSearcher.java b/lucene/suggest/src/java/org/apache/lucene/search/suggest/document/SuggestIndexSearcher.java index e46e73bb1aa9..0c88359029b4 100644 --- a/lucene/suggest/src/java/org/apache/lucene/search/suggest/document/SuggestIndexSearcher.java +++ b/lucene/suggest/src/java/org/apache/lucene/search/suggest/document/SuggestIndexSearcher.java @@ -22,6 +22,7 @@ import org.apache.lucene.search.BulkScorer; import org.apache.lucene.search.CollectionTerminatedException; import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.LeafCollector; import org.apache.lucene.search.Weight; /** @@ -67,14 +68,16 @@ public void suggest(CompletionQuery query, TopSuggestDocsCollector collector) th for (LeafReaderContext context : getIndexReader().leaves()) { BulkScorer scorer = weight.bulkScorer(context); if (scorer != null) { + LeafCollector leafCollector = collector.getLeafCollector(context); try { - scorer.score(collector.getLeafCollector(context), context.reader().getLiveDocs()); + scorer.score(leafCollector, context.reader().getLiveDocs()); } catch ( @SuppressWarnings("unused") CollectionTerminatedException e) { // collection was terminated prematurely // continue with the following leaf } + leafCollector.finish(); } } } diff --git a/lucene/suggest/src/java/org/apache/lucene/search/suggest/document/TopSuggestDocsCollector.java b/lucene/suggest/src/java/org/apache/lucene/search/suggest/document/TopSuggestDocsCollector.java index e8f66c87b3b7..3cfc8e2582df 100644 --- a/lucene/suggest/src/java/org/apache/lucene/search/suggest/document/TopSuggestDocsCollector.java +++ b/lucene/suggest/src/java/org/apache/lucene/search/suggest/document/TopSuggestDocsCollector.java @@ -100,12 +100,19 @@ public int getCountToCollect() { @Override protected void doSetNextReader(LeafReaderContext context) throws IOException { docBase = context.docBase; + } + + @Override + public void finish() throws IOException { if (seenSurfaceForms != null) { - seenSurfaceForms.clear(); // NOTE: this also clears the priorityQueue: for (SuggestScoreDoc hit : priorityQueue.getResults()) { pendingResults.add(hit); } + + // Deduplicate all hits: we already dedup'd efficiently within each segment by + // truncating the FST top paths search, but across segments there may still be dups: + seenSurfaceForms.clear(); } } @@ -136,15 +143,7 @@ public TopSuggestDocs get() throws IOException { SuggestScoreDoc[] suggestScoreDocs; if (seenSurfaceForms != null) { - // NOTE: this also clears the priorityQueue: - for (SuggestScoreDoc hit : priorityQueue.getResults()) { - pendingResults.add(hit); - } - - // Deduplicate all hits: we already dedup'd efficiently within each segment by - // truncating the FST top paths search, but across segments there may still be dups: - seenSurfaceForms.clear(); - + assert seenSurfaceForms.isEmpty(); // TODO: we could use a priority queue here to make cost O(N * log(num)) instead of O(N * // log(N)), where N = O(num * // numSegments), but typically numSegments is smallish and num is smallish so this won't diff --git a/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingCollector.java b/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingCollector.java index cf2c2732614d..af0df8cdc3f8 100644 --- a/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingCollector.java +++ b/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingCollector.java @@ -30,11 +30,12 @@ class AssertingCollector extends FilterCollector { private boolean weightSet = false; private int maxDoc = -1; private int previousLeafMaxDoc = 0; + boolean hasFinishedCollectingPreviousLeaf = true; /** Wrap the given collector in order to add assertions. */ - public static Collector wrap(Collector in) { + public static AssertingCollector wrap(Collector in) { if (in instanceof AssertingCollector) { - return in; + return (AssertingCollector) in; } return new AssertingCollector(in); } @@ -49,7 +50,9 @@ public LeafCollector getLeafCollector(LeafReaderContext context) throws IOExcept assert context.docBase >= previousLeafMaxDoc; previousLeafMaxDoc = context.docBase + context.reader().maxDoc(); + assert hasFinishedCollectingPreviousLeaf; final LeafCollector in = super.getLeafCollector(context); + hasFinishedCollectingPreviousLeaf = false; final int docBase = context.docBase; return new AssertingLeafCollector(in, 0, DocIdSetIterator.NO_MORE_DOCS) { @Override @@ -66,6 +69,12 @@ public void collect(int doc) throws IOException { super.collect(doc); maxDoc = docBase + doc; } + + @Override + public void finish() throws IOException { + hasFinishedCollectingPreviousLeaf = true; + super.finish(); + } }; } diff --git a/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingIndexSearcher.java b/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingIndexSearcher.java index 83abd7883091..f5fa29b14948 100644 --- a/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingIndexSearcher.java +++ b/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingIndexSearcher.java @@ -75,7 +75,9 @@ public Query rewrite(Query original) throws IOException { protected void search(List leaves, Weight weight, Collector collector) throws IOException { assert weight instanceof AssertingWeight; - super.search(leaves, weight, AssertingCollector.wrap(collector)); + AssertingCollector assertingCollector = AssertingCollector.wrap(collector); + super.search(leaves, weight, assertingCollector); + assert assertingCollector.hasFinishedCollectingPreviousLeaf; } @Override diff --git a/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingLeafCollector.java b/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingLeafCollector.java index 5c7801e51220..bcf3c4b1098b 100644 --- a/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingLeafCollector.java +++ b/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingLeafCollector.java @@ -30,6 +30,7 @@ class AssertingLeafCollector extends FilterLeafCollector { private Scorable scorer; private int lastCollected = -1; + private boolean finishCalled; AssertingLeafCollector(LeafCollector collector, int min, int max) { super(collector); @@ -57,4 +58,11 @@ public void collect(int doc) throws IOException { public DocIdSetIterator competitiveIterator() throws IOException { return in.competitiveIterator(); } + + @Override + public void finish() throws IOException { + assert finishCalled == false; + finishCalled = true; + super.finish(); + } }