From 837c512d7a11ded4ce7ee20315c544002fe0992e Mon Sep 17 00:00:00 2001 From: sravichandran Date: Sat, 25 Jul 2026 22:18:14 +0530 Subject: [PATCH] Fix over-allocation of result bitset in MemoryAccountingBitsetCollectorManager Size the reduce() result FixedBitSet to the highest matched doc + 1 instead of the largest maxDocEnd across collectors. The old sizing was inflated to the last visited leaf's docBase+maxDoc even when the query matched only a small subset, and under intra-segment concurrency it could equal the full index maxDoc. --- lucene/CHANGES.txt | 3 + .../MemoryAccountingBitsetCollector.java | 19 +++- ...emoryAccountingBitsetCollectorManager.java | 36 +++++--- .../TestMemoryAccountingBitsetCollector.java | 88 +++++++++++++++++++ 4 files changed, 131 insertions(+), 15 deletions(-) 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()); + } }