diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index d361021e414d..a2e96023e249 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -434,6 +434,9 @@ Bug Fixes
GlobalOrdinalsWithScoreCollector, fixing an intermittent failure caused by
non-associative float addition that #16378 missed. (Luca Cavanna)
+* GITHUB#16452: Fix over-allocation in MemoryAccountingBitsetCollectorManager.Result#bitSet, which
+ is now sized to highestMatchedDoc + 1 instead of the full searched range. (Sasilekha R)
+
Other
---------------------
* GITHUB#16266: Remove deprecated search(Query, Collector) calls in QueryUtils by replacing
diff --git a/lucene/misc/src/java/org/apache/lucene/misc/search/MemoryAccountingBitsetCollector.java b/lucene/misc/src/java/org/apache/lucene/misc/search/MemoryAccountingBitsetCollector.java
index 386921a231d7..f705f60bc24c 100644
--- a/lucene/misc/src/java/org/apache/lucene/misc/search/MemoryAccountingBitsetCollector.java
+++ b/lucene/misc/src/java/org/apache/lucene/misc/search/MemoryAccountingBitsetCollector.java
@@ -34,6 +34,12 @@ public class MemoryAccountingBitsetCollector extends SimpleCollector {
int minDocBase = Integer.MAX_VALUE;
int maxDocEnd = 0;
+ // Highest bit index set in bitSet, or -1 if no doc has been collected. Docs are collected in
+ // strictly ascending order: within a leaf by the Collector contract, and across leaves for a
+ // given collector because IndexSearcher sorts partitions within a slice by docBase and rejects
+ // multiple partitions of the same leaf sharing a slice. So this is simply the position written
+ // by the most recent collect() call.
+ int highestSetBit = -1;
public MemoryAccountingBitsetCollector(CollectorMemoryTracker tracker) {
this.tracker = tracker;
@@ -57,7 +63,14 @@ protected void doSetNextReader(LeafReaderContext context) throws IOException {
@Override
public void collect(int doc) {
- bitSet.set(docBase - minDocBase + doc);
+ int local = docBase - minDocBase + doc;
+ assert local > highestSetBit
+ : "collect() must receive docs in strictly ascending order; got local="
+ + local
+ + " after highestSetBit="
+ + highestSetBit;
+ bitSet.set(local);
+ highestSetBit = local;
}
@Override
@@ -69,7 +82,7 @@ int getMinDocBase() {
return minDocBase;
}
- int getMaxDocEnd() {
- return maxDocEnd;
+ int getHighestSetBit() {
+ return highestSetBit;
}
}
diff --git a/lucene/misc/src/java/org/apache/lucene/misc/search/MemoryAccountingBitsetCollectorManager.java b/lucene/misc/src/java/org/apache/lucene/misc/search/MemoryAccountingBitsetCollectorManager.java
index 85e8110843ff..223b85496de6 100644
--- a/lucene/misc/src/java/org/apache/lucene/misc/search/MemoryAccountingBitsetCollectorManager.java
+++ b/lucene/misc/src/java/org/apache/lucene/misc/search/MemoryAccountingBitsetCollectorManager.java
@@ -24,14 +24,22 @@
/**
* CollectorManager for MemoryAccountingBitsetCollector that supports concurrent search.
*
- * Creates multiple collectors for concurrent execution, each collector only allocates bitset for
- * slices it processes, then merges with proper offset in reduce().
+ *
Creates multiple collectors for concurrent execution; each collector only allocates a bitset
+ * for the slices it processes, and {@link #reduce} merges them into a single {@link Result} sized
+ * to the highest matched document across all collectors.
*/
public class MemoryAccountingBitsetCollectorManager
implements CollectorManager<
MemoryAccountingBitsetCollector, MemoryAccountingBitsetCollectorManager.Result> {
- /** The result of a search, containing the matched document IDs and total memory used. */
+ /**
+ * The result of a search, containing the matched document IDs and total memory used.
+ *
+ *
The returned {@link FixedBitSet} has length {@code highestMatchedDoc + 1}, or {@code 0} if
+ * no document matched; it is not padded to the searched index range. Callers probing document ids
+ * beyond the last match should first check {@link FixedBitSet#length()} before calling {@link
+ * FixedBitSet#get(int)}.
+ */
public record Result(FixedBitSet bitSet, long totalBytesUsed) {}
private final CollectorMemoryTracker tracker;
@@ -47,21 +55,25 @@ public MemoryAccountingBitsetCollector newCollector() {
@Override
public Result reduce(Collection collectors) {
- int globalMaxDocEnd = 0;
+ // Size the result to just cover the highest matched doc across all collectors. Each
+ // collector's maxDocEnd is inflated by doSetNextReader to the full leaf regardless of what
+ // actually matches, so keying off it can significantly over-allocate on selective queries or
+ // narrow intra-segment slices; use the actual high-water mark tracked at collect time.
+ int resultSize = 0;
for (MemoryAccountingBitsetCollector collector : collectors) {
- globalMaxDocEnd = Math.max(globalMaxDocEnd, collector.getMaxDocEnd());
+ int last = collector.getHighestSetBit();
+ if (last >= 0) {
+ resultSize = Math.max(resultSize, collector.getMinDocBase() + last + 1);
+ }
}
- // TODO: with intra-segment concurrency enabled, globalMaxDocEnd equals the full index maxDoc
- // even when only a portion of the index was searched, causing over-allocation of the result
- // bitset.
- FixedBitSet result = new FixedBitSet(globalMaxDocEnd);
+ FixedBitSet result = new FixedBitSet(resultSize);
tracker.updateBytes(result.ramBytesUsed());
for (MemoryAccountingBitsetCollector collector : collectors) {
- if (collector.bitSet != null && collector.bitSet.length() > 0) {
- int length = collector.getMaxDocEnd() - collector.getMinDocBase();
- FixedBitSet.orRange(collector.bitSet, 0, result, collector.getMinDocBase(), length);
+ int last = collector.getHighestSetBit();
+ if (last >= 0) {
+ FixedBitSet.orRange(collector.bitSet, 0, result, collector.getMinDocBase(), last + 1);
}
}
diff --git a/lucene/misc/src/test/org/apache/lucene/misc/search/TestMemoryAccountingBitsetCollector.java b/lucene/misc/src/test/org/apache/lucene/misc/search/TestMemoryAccountingBitsetCollector.java
index 7bbbe52102ac..c4c82b066e6f 100644
--- a/lucene/misc/src/test/org/apache/lucene/misc/search/TestMemoryAccountingBitsetCollector.java
+++ b/lucene/misc/src/test/org/apache/lucene/misc/search/TestMemoryAccountingBitsetCollector.java
@@ -17,13 +17,19 @@
package org.apache.lucene.misc.search;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.SortedDocValuesField;
import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.Term;
import org.apache.lucene.misc.CollectorMemoryTracker;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.MatchAllDocsQuery;
+import org.apache.lucene.search.TermQuery;
import org.apache.lucene.store.Directory;
import org.apache.lucene.tests.index.RandomIndexWriter;
import org.apache.lucene.tests.util.LuceneTestCase;
@@ -86,4 +92,86 @@ public void testCollectedResult() throws Exception {
// For collector with collecting only 1 doc, 80 bytes are required.
assertTrue(result.totalBytesUsed() >= 80);
}
+
+ public void testResultBitSetSizedToHighestMatchedDoc() throws Exception {
+ // Highly selective query: matches a single document early in the index. The result bitset
+ // must be sized tightly to (highestMatchedDoc + 1) rather than padded to the last visited
+ // leaf's docBase + maxDoc. Uses newSearcher() to get randomized executor/slicing coverage;
+ // the deterministic intra-segment case is exercised by the sister test below.
+ CollectorMemoryTracker tracker =
+ new CollectorMemoryTracker("testMemoryTracker", Long.MAX_VALUE);
+ MemoryAccountingBitsetCollectorManager bitsetCollectorManager =
+ new MemoryAccountingBitsetCollectorManager(tracker);
+
+ IndexSearcher searcher = newSearcher(reader);
+ MemoryAccountingBitsetCollectorManager.Result result =
+ searcher.search(new TermQuery(new Term("field", "5")), bitsetCollectorManager);
+
+ assertEquals(1, result.bitSet().cardinality());
+ int matchedDoc = result.bitSet().nextSetBit(0);
+ assertEquals(matchedDoc + 1, result.bitSet().length());
+ assertTrue(result.bitSet().length() < reader.maxDoc());
+ }
+
+ public void testResultBitSetSizedToHighestMatchedDocUnderIntraSegmentConcurrency()
+ throws Exception {
+ CollectorMemoryTracker tracker =
+ new CollectorMemoryTracker("testMemoryTracker", Long.MAX_VALUE);
+ MemoryAccountingBitsetCollectorManager bitsetCollectorManager =
+ new MemoryAccountingBitsetCollectorManager(tracker);
+
+ IndexSearcher searcher =
+ new IndexSearcher(reader, Runnable::run) {
+ @Override
+ protected LeafSlice[] slices(List leaves) {
+ // Split each leaf into two partitions, each in its own slice, to force
+ // intra-segment concurrency.
+ List slices = new ArrayList<>();
+ for (LeafReaderContext ctx : leaves) {
+ int maxDoc = ctx.reader().maxDoc();
+ if (maxDoc <= 1) {
+ slices.add(
+ new LeafSlice(
+ Collections.singletonList(
+ LeafReaderContextPartition.createForEntireSegment(ctx))));
+ } else {
+ int mid = maxDoc / 2;
+ slices.add(
+ new LeafSlice(
+ Collections.singletonList(
+ LeafReaderContextPartition.createFromAndTo(ctx, 0, mid))));
+ slices.add(
+ new LeafSlice(
+ Collections.singletonList(
+ LeafReaderContextPartition.createFromAndTo(ctx, mid, maxDoc))));
+ }
+ }
+ return slices.toArray(LeafSlice[]::new);
+ }
+ };
+
+ // Use a mid-index match so any given leaf's slice split lands the match in a non-trivial
+ // partition, and the highest-set bit differs from the other tests.
+ MemoryAccountingBitsetCollectorManager.Result result =
+ searcher.search(new TermQuery(new Term("field", "500")), bitsetCollectorManager);
+
+ assertEquals(1, result.bitSet().cardinality());
+ int matchedDoc = result.bitSet().nextSetBit(0);
+ assertEquals(matchedDoc + 1, result.bitSet().length());
+ assertTrue(result.bitSet().length() < reader.maxDoc());
+ }
+
+ public void testResultBitSetEmptyOnNoMatches() throws Exception {
+ CollectorMemoryTracker tracker =
+ new CollectorMemoryTracker("testMemoryTracker", Long.MAX_VALUE);
+ MemoryAccountingBitsetCollectorManager bitsetCollectorManager =
+ new MemoryAccountingBitsetCollectorManager(tracker);
+
+ IndexSearcher searcher = newSearcher(reader);
+ MemoryAccountingBitsetCollectorManager.Result result =
+ searcher.search(new TermQuery(new Term("field", "does-not-exist")), bitsetCollectorManager);
+
+ assertEquals(0, result.bitSet().cardinality());
+ assertEquals(0, result.bitSet().length());
+ }
}