Fix over-allocation of result bitset in MemoryAccountingBitsetCollectorManager

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

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;
   }
 }

lucene/misc/src/java/org/apache/lucene/misc/search/MemoryAccountingBitsetCollectorManager.java

@@ -24,14 +24,22 @@
 /**
  * CollectorManager for MemoryAccountingBitsetCollector that supports concurrent search.
  *
- * <p>Creates multiple collectors for concurrent execution, each collector only allocates bitset for
- * slices it processes, then merges with proper offset in reduce().
+ * <p>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.
+   *
+   * <p>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<MemoryAccountingBitsetCollector> 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);
       }
     }
 

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<LeafReaderContext> leaves) {
+            // Split each leaf into two partitions, each in its own slice, to force
+            // intra-segment concurrency.
+            List<LeafSlice> 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());
+  }
 }