Optimize AbstractKnnVectorQuery#createBitSet with intoBitset (#14674)

Author: gf2121

lucene/CHANGES.txt

@@ -65,6 +65,8 @@ Optimizations
 
 * GITHUB#14709: Speed up TermQuery by Scorer#nextDocsAndScores. (Guo Feng)
 
+* GITHUB#14674: Optimize AbstractKnnVectorQuery#createBitSet with intoBitset. (Guo Feng)
+
 Bug Fixes
 ---------------------
 * GITHUB#14654: ValueSource.fromDoubleValuesSource(dvs).getSortField() would throw errors when

lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java

@@ -37,6 +37,7 @@
 import org.apache.lucene.util.BitSet;
 import org.apache.lucene.util.BitSetIterator;
 import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.FixedBitSet;
 
 /**
  * Uses {@link KnnVectorsReader#search} to perform nearest neighbour search.
@@ -171,15 +172,25 @@ private BitSet createBitSet(DocIdSetIterator iterator, Bits liveDocs, int maxDoc
       // If we already have a BitSet and no deletions, reuse the BitSet
       return bitSetIterator.getBitSet();
     } else {
-      // Create a new BitSet from matching and live docs
-      FilteredDocIdSetIterator filterIterator =
-          new FilteredDocIdSetIterator(iterator) {
-            @Override
-            protected boolean match(int doc) {
-              return liveDocs == null || liveDocs.get(doc);
-            }
-          };
-      return BitSet.of(filterIterator, maxDoc);
+      int threshold = maxDoc >> 7; // same as BitSet#of
+      if (iterator.cost() >= threshold) {
+        // take advantage of Disi#intoBitset and Bits#applyMask
+        FixedBitSet bitSet = new FixedBitSet(maxDoc);
+        bitSet.or(iterator);
+        if (liveDocs != null) {
+          liveDocs.applyMask(bitSet, 0);
+        }
+        return bitSet;
+      } else {
+        FilteredDocIdSetIterator filterIterator =
+            new FilteredDocIdSetIterator(iterator) {
+              @Override
+              protected boolean match(int doc) {
+                return liveDocs == null || liveDocs.get(doc);
+              }
+            };
+        return BitSet.of(filterIterator, maxDoc); // create a sparse bitset
+      }
     }
   }
 

lucene/core/src/java/org/apache/lucene/search/AbstractVectorSimilarityQuery.java

@@ -29,6 +29,7 @@
 import org.apache.lucene.util.BitSet;
 import org.apache.lucene.util.BitSetIterator;
 import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.FixedBitSet;
 
 /**
  * Search for all (approximate) vectors above a similarity threshold.
@@ -142,14 +143,27 @@ public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOExcepti
             acceptDocs = bitSetIterator.getBitSet();
           } else {
             // Else collect all matching docs
-            FilteredDocIdSetIterator filtered =
-                new FilteredDocIdSetIterator(scorer.iterator()) {
-                  @Override
-                  protected boolean match(int doc) {
-                    return liveDocs == null || liveDocs.get(doc);
-                  }
-                };
-            acceptDocs = BitSet.of(filtered, leafReader.maxDoc());
+            DocIdSetIterator iterator = scorer.iterator();
+            final int maxDoc = leafReader.maxDoc();
+            int threshold = maxDoc >> 7; // same as BitSet#of
+            if (iterator.cost() >= threshold) {
+              // take advantage of Disi#intoBitset and Bits#applyMask
+              FixedBitSet bitSet = new FixedBitSet(maxDoc);
+              bitSet.or(iterator);
+              if (liveDocs != null) {
+                liveDocs.applyMask(bitSet, 0);
+              }
+              acceptDocs = bitSet;
+            } else {
+              FilteredDocIdSetIterator filterIterator =
+                  new FilteredDocIdSetIterator(iterator) {
+                    @Override
+                    protected boolean match(int doc) {
+                      return liveDocs == null || liveDocs.get(doc);
+                    }
+                  };
+              acceptDocs = BitSet.of(filterIterator, maxDoc); // create a sparse bitset
+            }
           }
 
           int cardinality = acceptDocs.cardinality();