You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by ja...@apache.org on 2026/06/15 13:21:04 UTC
(lucene) branch main updated: Add DiversifiedTopDocsCollectorManager to enable concurrent search (#16216)
This is an automated email from the ASF dual-hosted git repository.
javanna pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/lucene.git
The following commit(s) were added to refs/heads/main by this push:
new 0a598f6b3b5 Add DiversifiedTopDocsCollectorManager to enable concurrent search (#16216)
0a598f6b3b5 is described below
commit 0a598f6b3b5b1eb04f7a663c80352ff934f13b81
Author: Luca Cavanna <ja...@apache.org>
AuthorDate: Mon Jun 15 15:20:58 2026 +0200
Add DiversifiedTopDocsCollectorManager to enable concurrent search (#16216)
Introduces DiversifiedTopDocsCollectorManager, an abstract CollectorManager
whose reduce() merges per-slice results while re-applying the diversity
constraint globally. Updates TestDiversifiedTopDocsCollector to use it,
removing usages of the deprecated IndexSearcher#search(Query, Collector).
Relates to #12892
---
lucene/CHANGES.txt | 4 +
.../search/DiversifiedTopDocsCollectorManager.java | 93 +++++++++++++++++
.../search/TestDiversifiedTopDocsCollector.java | 114 +++++++++++++--------
3 files changed, 168 insertions(+), 43 deletions(-)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index a961549de83..38b7958904a 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -402,6 +402,10 @@ Improvements
* GITHUB#16217: Introduce DistinctValuesCollectorManager to parallelize search when using DistinctValuesCollector. (Luca Cavanna)
+* GITHUB#16216: Add DiversifiedTopDocsCollectorManager to enable concurrent search with
+ DiversifiedTopDocsCollector, removing usages of the deprecated IndexSearcher#search(Query, Collector).
+ (Luca Cavanna)
+
Optimizations
---------------------
* GITHUB#16222: MultiTermQuery constant-score wrapper now defers term collection to ScorerSupplier#get()
diff --git a/lucene/misc/src/java/org/apache/lucene/misc/search/DiversifiedTopDocsCollectorManager.java b/lucene/misc/src/java/org/apache/lucene/misc/search/DiversifiedTopDocsCollectorManager.java
new file mode 100644
index 00000000000..a059e59df81
--- /dev/null
+++ b/lucene/misc/src/java/org/apache/lucene/misc/search/DiversifiedTopDocsCollectorManager.java
@@ -0,0 +1,93 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.misc.search;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import org.apache.lucene.internal.hppc.LongIntHashMap;
+import org.apache.lucene.misc.search.DiversifiedTopDocsCollector.ScoreDocKey;
+import org.apache.lucene.search.CollectorManager;
+import org.apache.lucene.search.ScoreDoc;
+import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.search.TotalHits;
+
+/**
+ * A {@link CollectorManager} for {@link DiversifiedTopDocsCollector}. Subclasses implement {@link
+ * #newCollector()} to supply concrete collector instances; this class provides the {@link
+ * #reduce(Collection)} that merges per-slice results while preserving the diversity constraint
+ * across the whole result set.
+ *
+ * @lucene.experimental
+ */
+public abstract class DiversifiedTopDocsCollectorManager<C extends DiversifiedTopDocsCollector>
+ implements CollectorManager<C, TopDocs> {
+
+ private final int numHits;
+ private final int maxHitsPerKey;
+
+ /**
+ * @param numHits maximum number of results to return
+ * @param maxHitsPerKey maximum number of results per key allowed in the final results
+ */
+ protected DiversifiedTopDocsCollectorManager(int numHits, int maxHitsPerKey) {
+ this.numHits = numHits;
+ this.maxHitsPerKey = maxHitsPerKey;
+ }
+
+ /**
+ * Merges per-slice results into a single {@link TopDocs} while re-applying the diversity
+ * constraint globally. Each per-slice collector may contain up to {@code numHits} results, but
+ * the same key can appear in multiple slices; {@code reduce} enforces {@code maxHitsPerKey}
+ * across all slices.
+ */
+ @Override
+ public TopDocs reduce(Collection<C> collectors) throws IOException {
+ long totalHits = 0;
+ List<ScoreDocKey> allDocs = new ArrayList<>();
+ for (C collector : collectors) {
+ TopDocs topDocs = collector.topDocs();
+ totalHits += topDocs.totalHits.value();
+ for (ScoreDoc sd : topDocs.scoreDocs) {
+ allDocs.add((ScoreDocKey) sd);
+ }
+ }
+ // Sort descending by score; break ties ascending by doc for determinism
+ allDocs.sort(
+ (a, b) -> {
+ int cmp = Float.compare(b.score, a.score);
+ return cmp != 0 ? cmp : Integer.compare(a.doc, b.doc);
+ });
+ // Greedy pick: take up to numHits while honouring maxHitsPerKey per key
+ LongIntHashMap keyCounts = new LongIntHashMap();
+ List<ScoreDoc> result = new ArrayList<>();
+ for (ScoreDocKey sdk : allDocs) {
+ assert sdk.key != null;
+ int count = keyCounts.getOrDefault(sdk.key, 0);
+ if (count < maxHitsPerKey) {
+ result.add(sdk);
+ keyCounts.put(sdk.key, count + 1);
+ if (result.size() == numHits) {
+ break;
+ }
+ }
+ }
+ return new TopDocs(
+ new TotalHits(totalHits, TotalHits.Relation.EQUAL_TO), result.toArray(ScoreDoc[]::new));
+ }
+}
diff --git a/lucene/misc/src/test/org/apache/lucene/misc/search/TestDiversifiedTopDocsCollector.java b/lucene/misc/src/test/org/apache/lucene/misc/search/TestDiversifiedTopDocsCollector.java
index a5f5e9e7f1e..4cf850f0d52 100644
--- a/lucene/misc/src/test/org/apache/lucene/misc/search/TestDiversifiedTopDocsCollector.java
+++ b/lucene/misc/src/test/org/apache/lucene/misc/search/TestDiversifiedTopDocsCollector.java
@@ -17,6 +17,7 @@
package org.apache.lucene.misc.search;
import java.io.IOException;
+import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import org.apache.lucene.document.Document;
@@ -29,11 +30,13 @@ import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.MultiDocValues;
import org.apache.lucene.index.NumericDocValues;
+import org.apache.lucene.index.ReaderUtil;
import org.apache.lucene.index.SortedDocValues;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.BooleanClause;
import org.apache.lucene.search.BooleanClause.Occur;
import org.apache.lucene.search.BooleanQuery;
+import org.apache.lucene.search.BulkScorer;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.Explanation;
import org.apache.lucene.search.IndexSearcher;
@@ -76,9 +79,8 @@ public class TestDiversifiedTopDocsCollector extends LuceneTestCase {
// any one artist.
int requiredMaxHitsPerArtist = 2;
int numberOfTracksOnCompilation = 10;
- DiversifiedTopDocsCollector tdc =
- doDiversifiedSearch(numberOfTracksOnCompilation, requiredMaxHitsPerArtist);
- ScoreDoc[] sd = tdc.topDocs(0).scoreDocs;
+ TopDocs results = doDiversifiedSearch(numberOfTracksOnCompilation, requiredMaxHitsPerArtist);
+ ScoreDoc[] sd = results.scoreDocs;
assertEquals(numberOfTracksOnCompilation, sd.length);
assertTrue(getMaxNumRecordsPerArtist(sd) <= requiredMaxHitsPerArtist);
}
@@ -89,59 +91,77 @@ public class TestDiversifiedTopDocsCollector extends LuceneTestCase {
int requiredMaxHitsPerArtist = 1;
// Volume 2 of our hits compilation - start at position 10
- DiversifiedTopDocsCollector tdc =
+ TopDocs td =
doDiversifiedSearch(
numberOfTracksPerCompilation * numberOfCompilations, requiredMaxHitsPerArtist);
ScoreDoc[] volume2 =
- tdc.topDocs(numberOfTracksPerCompilation, numberOfTracksPerCompilation).scoreDocs;
+ Arrays.copyOfRange(
+ td.scoreDocs,
+ numberOfTracksPerCompilation,
+ numberOfTracksPerCompilation * numberOfCompilations);
assertEquals(numberOfTracksPerCompilation, volume2.length);
assertTrue(getMaxNumRecordsPerArtist(volume2) <= requiredMaxHitsPerArtist);
}
public void testInvalidArguments() throws Exception {
int numResults = 5;
- DiversifiedTopDocsCollector tdc = doDiversifiedSearch(numResults, 15);
+ // Create an empty collector to exercise topDocs() argument validation
+ DiversifiedTopDocsCollector tdc =
+ new HashedDocValuesDiversifiedCollector(numResults, 15, "artist");
- // start < 0
+ // start < 0 throws
IllegalArgumentException expected =
- expectThrows(
- IllegalArgumentException.class,
- () -> {
- tdc.topDocs(-1);
- });
-
+ expectThrows(IllegalArgumentException.class, () -> tdc.topDocs(-1));
assertEquals(
- "Expected value of starting position is between 0 and 5, got -1", expected.getMessage());
+ "Expected value of starting position is between 0 and 0, got -1", expected.getMessage());
- // start > pq.size()
+ // start > pq.size() (pq is empty; any positive start is out of range)
assertEquals(0, tdc.topDocs(numResults + 1).scoreDocs.length);
- // start == pq.size()
- assertEquals(0, tdc.topDocs(numResults).scoreDocs.length);
+ // start == pq.size() == 0
+ assertEquals(0, tdc.topDocs(0).scoreDocs.length);
// howMany < 0
- expected =
- expectThrows(
- IllegalArgumentException.class,
- () -> {
- tdc.topDocs(0, -1);
- });
-
+ expected = expectThrows(IllegalArgumentException.class, () -> tdc.topDocs(0, -1));
assertEquals(
"Number of hits requested must be greater than 0 but value was -1", expected.getMessage());
// howMany == 0
assertEquals(0, tdc.topDocs(0, 0).scoreDocs.length);
+
+ // Verify start == pq.size() returns empty with a non-empty collector
+ HashedDocValuesDiversifiedCollector populated =
+ new HashedDocValuesDiversifiedCollector(numResults, 15, "artist");
+ Weight w = searcher.createWeight(searcher.rewrite(getTestQuery()), ScoreMode.COMPLETE, 1f);
+ for (LeafReaderContext ctx : reader.leaves()) {
+ LeafCollector lc = populated.getLeafCollector(ctx);
+ BulkScorer bs = w.bulkScorer(ctx);
+ if (bs != null) {
+ bs.score(lc, ctx.reader().getLiveDocs(), 0, DocIdSetIterator.NO_MORE_DOCS);
+ }
+ }
+ int collectedCount = populated.topDocs().scoreDocs.length;
+ assertTrue(collectedCount > 0);
+ assertEquals(0, populated.topDocs(collectedCount).scoreDocs.length);
}
// Diversifying collector that looks up de-dup keys using SortedDocValues
// from a top-level Reader
private static final class DocValuesDiversifiedCollector extends DiversifiedTopDocsCollector {
- private final SortedDocValues sdv;
+ private SortedDocValues sdv;
- public DocValuesDiversifiedCollector(int size, int maxHitsPerKey, SortedDocValues sdv) {
+ public DocValuesDiversifiedCollector(int size, int maxHitsPerKey) {
super(size, maxHitsPerKey);
- this.sdv = sdv;
+ }
+
+ @Override
+ public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
+ if (sdv == null) {
+ sdv =
+ MultiDocValues.getSortedValues(
+ ReaderUtil.getTopLevelContext(context).reader(), "artist");
+ }
+ return super.getLeafCollector(context);
}
@Override
@@ -322,7 +342,6 @@ public class TestDiversifiedTopDocsCollector extends LuceneTestCase {
private Directory dir;
private IndexReader reader;
private IndexSearcher searcher;
- private SortedDocValues artistDocValues;
static class Record {
String year;
@@ -356,8 +375,7 @@ public class TestDiversifiedTopDocsCollector extends LuceneTestCase {
}
}
- private DiversifiedTopDocsCollector doDiversifiedSearch(int numResults, int maxResultsPerArtist)
- throws IOException {
+ private TopDocs doDiversifiedSearch(int numResults, int maxResultsPerArtist) throws IOException {
// Alternate between implementations used for key lookups
if (random().nextBoolean()) {
// Faster key lookup but with potential for collisions on larger datasets
@@ -368,20 +386,31 @@ public class TestDiversifiedTopDocsCollector extends LuceneTestCase {
}
}
- private DiversifiedTopDocsCollector doFuzzyDiversifiedSearch(
- int numResults, int maxResultsPerArtist) throws IOException {
- DiversifiedTopDocsCollector tdc =
- new HashedDocValuesDiversifiedCollector(numResults, maxResultsPerArtist, "artist");
- searcher.search(getTestQuery(), tdc);
- return tdc;
+ private TopDocs doFuzzyDiversifiedSearch(int numResults, int maxResultsPerArtist)
+ throws IOException {
+ return searcher.search(
+ getTestQuery(),
+ new DiversifiedTopDocsCollectorManager<HashedDocValuesDiversifiedCollector>(
+ numResults, maxResultsPerArtist) {
+ @Override
+ public HashedDocValuesDiversifiedCollector newCollector() {
+ return new HashedDocValuesDiversifiedCollector(
+ numResults, maxResultsPerArtist, "artist");
+ }
+ });
}
- private DiversifiedTopDocsCollector doAccurateDiversifiedSearch(
- int numResults, int maxResultsPerArtist) throws IOException {
- DiversifiedTopDocsCollector tdc =
- new DocValuesDiversifiedCollector(numResults, maxResultsPerArtist, artistDocValues);
- searcher.search(getTestQuery(), tdc);
- return tdc;
+ private TopDocs doAccurateDiversifiedSearch(int numResults, int maxResultsPerArtist)
+ throws IOException {
+ return searcher.search(
+ getTestQuery(),
+ new DiversifiedTopDocsCollectorManager<DocValuesDiversifiedCollector>(
+ numResults, maxResultsPerArtist) {
+ @Override
+ public DocValuesDiversifiedCollector newCollector() {
+ return new DocValuesDiversifiedCollector(numResults, maxResultsPerArtist);
+ }
+ });
}
private Query getTestQuery() {
@@ -440,7 +469,6 @@ public class TestDiversifiedTopDocsCollector extends LuceneTestCase {
reader = writer.getReader();
writer.close();
searcher = newSearcher(reader);
- artistDocValues = MultiDocValues.getSortedValues(reader, "artist");
}
@Override