Add AcceptDocs abstraction for accepted KNN docs

lucene/CHANGES.txt

@@ -128,6 +128,10 @@ API Changes
 * GITHUB#14978: Add a bulk scoring interface to RandomVectorScorer
   (Trevor McCulloch, Chris Hegarty)
 
+* GITHUB#15011: LeafReader#searchNearestVectors now accepts an AcceptDocs
+  instance instead of a Bits instance to identify document IDs to filter.
+  (Shubham Chaudhary, Adrien Grand)
+
 New Features
 ---------------------
 * GITHUB#15015: MultiIndexMergeScheduler: a production multi-tenant merge scheduler (Shawn Yarbrough)

lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene90/Lucene90HnswVectorsReader.java

@@ -35,6 +35,7 @@
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.internal.hppc.IntObjectHashMap;
+import org.apache.lucene.search.AcceptDocs;
 import org.apache.lucene.search.KnnCollector;
 import org.apache.lucene.search.VectorScorer;
 import org.apache.lucene.store.ChecksumIndexInput;
@@ -242,7 +243,7 @@ public ByteVectorValues getByteVectorValues(String field) {
   }
 
   @Override
-  public void search(String field, float[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     final FieldEntry fieldEntry = getFieldEntry(field);
     if (fieldEntry.size() == 0) {
@@ -260,7 +261,7 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
             vectorValues,
             fieldEntry.similarityFunction,
             getGraphValues(fieldEntry),
-            getAcceptOrds(acceptDocs, fieldEntry),
+            getAcceptOrds(acceptDocs.bits(), fieldEntry),
             knnCollector.visitLimit(),
             random);
     knnCollector.incVisitedCount(results.visitedCount());
@@ -273,7 +274,7 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
   }
 
   @Override
-  public void search(String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     throw new UnsupportedOperationException();
   }

lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene91/Lucene91HnswVectorsReader.java

@@ -37,6 +37,7 @@
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.internal.hppc.IntObjectHashMap;
+import org.apache.lucene.search.AcceptDocs;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.KnnCollector;
 import org.apache.lucene.search.VectorScorer;
@@ -238,7 +239,7 @@ public ByteVectorValues getByteVectorValues(String field) throws IOException {
   }
 
   @Override
-  public void search(String field, float[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     final FieldEntry fieldEntry = getFieldEntry(field);
     if (fieldEntry.size() == 0) {
@@ -253,11 +254,11 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
         scorer,
         new OrdinalTranslatedKnnCollector(knnCollector, vectorValues::ordToDoc),
         getGraph(fieldEntry),
-        getAcceptOrds(acceptDocs, fieldEntry));
+        getAcceptOrds(acceptDocs.bits(), fieldEntry));
   }
 
   @Override
-  public void search(String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     throw new UnsupportedOperationException();
   }

lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene92/Lucene92HnswVectorsReader.java

@@ -36,11 +36,11 @@
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.internal.hppc.IntObjectHashMap;
+import org.apache.lucene.search.AcceptDocs;
 import org.apache.lucene.search.KnnCollector;
 import org.apache.lucene.store.ChecksumIndexInput;
 import org.apache.lucene.store.DataInput;
 import org.apache.lucene.store.IndexInput;
-import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.IOUtils;
 import org.apache.lucene.util.hnsw.HnswGraph;
 import org.apache.lucene.util.hnsw.HnswGraphSearcher;
@@ -236,7 +236,7 @@ public ByteVectorValues getByteVectorValues(String field) {
   }
 
   @Override
-  public void search(String field, float[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     final FieldEntry fieldEntry = getFieldEntry(field);
     if (fieldEntry.size() == 0) {
@@ -251,11 +251,11 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
         scorer,
         new OrdinalTranslatedKnnCollector(knnCollector, vectorValues::ordToDoc),
         getGraph(fieldEntry),
-        vectorValues.getAcceptOrds(acceptDocs));
+        vectorValues.getAcceptOrds(acceptDocs.bits()));
   }
 
   @Override
-  public void search(String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     throw new UnsupportedOperationException();
   }

lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene94/Lucene94HnswVectorsReader.java

@@ -37,11 +37,11 @@
 import org.apache.lucene.index.VectorEncoding;
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.internal.hppc.IntObjectHashMap;
+import org.apache.lucene.search.AcceptDocs;
 import org.apache.lucene.search.KnnCollector;
 import org.apache.lucene.store.ChecksumIndexInput;
 import org.apache.lucene.store.DataInput;
 import org.apache.lucene.store.IndexInput;
-import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.IOUtils;
 import org.apache.lucene.util.hnsw.HnswGraph;
 import org.apache.lucene.util.hnsw.HnswGraphSearcher;
@@ -270,7 +270,7 @@ public ByteVectorValues getByteVectorValues(String field) throws IOException {
   }
 
   @Override
-  public void search(String field, float[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     final FieldEntry fieldEntry = getFieldEntry(field, VectorEncoding.FLOAT32);
     if (fieldEntry.size() == 0 || knnCollector.k() == 0) {
@@ -285,11 +285,11 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
         scorer,
         new OrdinalTranslatedKnnCollector(knnCollector, vectorValues::ordToDoc),
         getGraph(fieldEntry),
-        vectorValues.getAcceptOrds(acceptDocs));
+        vectorValues.getAcceptOrds(acceptDocs.bits()));
   }
 
   @Override
-  public void search(String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     final FieldEntry fieldEntry = getFieldEntry(field, VectorEncoding.BYTE);
     if (fieldEntry.size() == 0 || knnCollector.k() == 0) {
@@ -304,7 +304,7 @@ public void search(String field, byte[] target, KnnCollector knnCollector, Bits
         scorer,
         new OrdinalTranslatedKnnCollector(knnCollector, vectorValues::ordToDoc),
         getGraph(fieldEntry),
-        vectorValues.getAcceptOrds(acceptDocs));
+        vectorValues.getAcceptOrds(acceptDocs.bits()));
   }
 
   private HnswGraph getGraph(FieldEntry entry) throws IOException {

lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene95/Lucene95HnswVectorsReader.java

@@ -41,13 +41,13 @@
 import org.apache.lucene.index.VectorEncoding;
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.internal.hppc.IntObjectHashMap;
+import org.apache.lucene.search.AcceptDocs;
 import org.apache.lucene.search.KnnCollector;
 import org.apache.lucene.store.ChecksumIndexInput;
 import org.apache.lucene.store.DataInput;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.store.RandomAccessInput;
 import org.apache.lucene.util.ArrayUtil;
-import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.IOUtils;
 import org.apache.lucene.util.hnsw.HnswGraph;
 import org.apache.lucene.util.hnsw.HnswGraphSearcher;
@@ -290,7 +290,7 @@ public ByteVectorValues getByteVectorValues(String field) throws IOException {
   }
 
   @Override
-  public void search(String field, float[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     final FieldEntry fieldEntry = getFieldEntry(field, VectorEncoding.FLOAT32);
     if (fieldEntry.size() == 0 || knnCollector.k() == 0) {
@@ -314,11 +314,11 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
         scorer,
         new OrdinalTranslatedKnnCollector(knnCollector, vectorValues::ordToDoc),
         getGraph(fieldEntry),
-        vectorValues.getAcceptOrds(acceptDocs));
+        vectorValues.getAcceptOrds(acceptDocs.bits()));
   }
 
   @Override
-  public void search(String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     final FieldEntry fieldEntry = getFieldEntry(field, VectorEncoding.BYTE);
     if (fieldEntry.size() == 0 || knnCollector.k() == 0) {
@@ -342,7 +342,7 @@ public void search(String field, byte[] target, KnnCollector knnCollector, Bits
         scorer,
         new OrdinalTranslatedKnnCollector(knnCollector, vectorValues::ordToDoc),
         getGraph(fieldEntry),
-        vectorValues.getAcceptOrds(acceptDocs));
+        vectorValues.getAcceptOrds(acceptDocs.bits()));
   }
 
   /** Get knn graph values; used for testing */

lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextKnnVectorsReader.java

@@ -37,14 +37,14 @@
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.internal.hppc.IntObjectHashMap;
+import org.apache.lucene.search.AcceptDocs;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.KnnCollector;
 import org.apache.lucene.search.VectorScorer;
 import org.apache.lucene.store.BufferedChecksumIndexInput;
 import org.apache.lucene.store.ChecksumIndexInput;
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexInput;
-import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.BytesRefBuilder;
 import org.apache.lucene.util.IOUtils;
@@ -178,7 +178,7 @@ public ByteVectorValues getByteVectorValues(String field) throws IOException {
   }
 
   @Override
-  public void search(String field, float[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     FloatVectorValues values = getFloatVectorValues(field);
     if (target.length != values.dimension()) {
@@ -192,7 +192,7 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
     VectorSimilarityFunction vectorSimilarity = info.getVectorSimilarityFunction();
     for (int ord = 0; ord < values.size(); ord++) {
       int doc = values.ordToDoc(ord);
-      if (acceptDocs != null && acceptDocs.get(doc) == false) {
+      if (acceptDocs.bits() != null && acceptDocs.bits().get(doc) == false) {
         continue;
       }
 
@@ -208,7 +208,7 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
   }
 
   @Override
-  public void search(String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     ByteVectorValues values = getByteVectorValues(field);
     if (target.length != values.dimension()) {
@@ -223,7 +223,7 @@ public void search(String field, byte[] target, KnnCollector knnCollector, Bits
 
     for (int ord = 0; ord < values.size(); ord++) {
       int doc = values.ordToDoc(ord);
-      if (acceptDocs != null && acceptDocs.get(doc) == false) {
+      if (acceptDocs.bits() != null && acceptDocs.bits().get(doc) == false) {
         continue;
       }
 

lucene/codecs/src/test/org/apache/lucene/codecs/bitvectors/TestHnswBitVectorsFormat.java

@@ -33,6 +33,7 @@
 import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.StoredFields;
 import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.AcceptDocs;
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.TopKnnCollector;
 import org.apache.lucene.store.Directory;
@@ -84,7 +85,8 @@ public void testIndexAndSearchBitVectors() throws IOException {
       try (IndexReader reader = DirectoryReader.open(w)) {
         LeafReader r = getOnlyLeafReader(reader);
         TopKnnCollector collector = new TopKnnCollector(3, Integer.MAX_VALUE);
-        r.searchNearestVectors("v1", vectors[0], collector, null);
+        r.searchNearestVectors(
+            "v1", vectors[0], collector, AcceptDocs.fromLiveDocs(null, r.maxDoc()));
         TopDocs topDocs = collector.topDocs();
         assertEquals(3, topDocs.scoreDocs.length);
 

lucene/core/src/java/org/apache/lucene/codecs/KnnVectorsFormat.java

@@ -25,8 +25,8 @@
 import org.apache.lucene.index.FloatVectorValues;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
+import org.apache.lucene.search.AcceptDocs;
 import org.apache.lucene.search.KnnCollector;
-import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.NamedSPILoader;
 
 /**
@@ -140,13 +140,13 @@ public ByteVectorValues getByteVectorValues(String field) {
 
             @Override
             public void search(
-                String field, float[] target, KnnCollector knnCollector, Bits acceptDocs) {
+                String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) {
               throw new UnsupportedOperationException();
             }
 
             @Override
             public void search(
-                String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs) {
+                String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) {
               throw new UnsupportedOperationException();
             }
 

lucene/core/src/java/org/apache/lucene/codecs/KnnVectorsReader.java

@@ -25,6 +25,7 @@
 import org.apache.lucene.index.ByteVectorValues;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.search.AcceptDocs;
 import org.apache.lucene.search.KnnCollector;
 import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.search.TopDocs;
@@ -88,7 +89,8 @@ protected KnnVectorsReader() {}
    *     if they are all allowed to match.
    */
   public abstract void search(
-      String field, float[] target, KnnCollector knnCollector, Bits acceptDocs) throws IOException;
+      String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
+      throws IOException;
 
   /**
    * Return the k nearest neighbor documents as determined by comparison of their vector values for
@@ -116,7 +118,8 @@ public abstract void search(
    *     if they are all allowed to match.
    */
   public abstract void search(
-      String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs) throws IOException;
+      String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
+      throws IOException;
 
   /**
    * Returns an instance optimized for merging. This instance may only be consumed in the thread

lucene/core/src/java/org/apache/lucene/codecs/hnsw/FlatVectorsReader.java

@@ -19,9 +19,9 @@
 
 import java.io.IOException;
 import org.apache.lucene.codecs.KnnVectorsReader;
+import org.apache.lucene.search.AcceptDocs;
 import org.apache.lucene.search.KnnCollector;
 import org.apache.lucene.util.Accountable;
-import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.hnsw.RandomVectorScorer;
 
 /**
@@ -56,13 +56,13 @@ public FlatVectorsScorer getFlatVectorScorer() {
   }
 
   @Override
-  public void search(String field, float[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     // don't scan stored field data. If we didn't index it, produce no search results
   }
 
   @Override
-  public void search(String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     // don't scan stored field data. If we didn't index it, produce no search results
   }

lucene/core/src/java/org/apache/lucene/codecs/lucene102/Lucene102BinaryQuantizedVectorsReader.java

@@ -38,6 +38,7 @@
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.VectorEncoding;
 import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.AcceptDocs;
 import org.apache.lucene.search.KnnCollector;
 import org.apache.lucene.search.VectorScorer;
 import org.apache.lucene.store.ChecksumIndexInput;
@@ -223,20 +224,20 @@ public ByteVectorValues getByteVectorValues(String field) throws IOException {
   }
 
   @Override
-  public void search(String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     rawVectorsReader.search(field, target, knnCollector, acceptDocs);
   }
 
   @Override
-  public void search(String field, float[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     if (knnCollector.k() == 0) return;
     final RandomVectorScorer scorer = getRandomVectorScorer(field, target);
     if (scorer == null) return;
     OrdinalTranslatedKnnCollector collector =
         new OrdinalTranslatedKnnCollector(knnCollector, scorer::ordToDoc);
-    Bits acceptedOrds = scorer.getAcceptOrds(acceptDocs);
+    Bits acceptedOrds = scorer.getAcceptOrds(acceptDocs.bits());
     for (int i = 0; i < scorer.maxOrd(); i++) {
       if (acceptedOrds == null || acceptedOrds.get(i)) {
         collector.collect(i, scorer.score(i));