Add new VectorScorer interface to vector value iterators (apache#13181)

With quantized vectors, and with current vectors, we separate out the "scoring" vs. "iteration", requiring the user to always iterate the raw vectors and provide their own similarity function.

While this is flexible, it creates frustration in:

 - Just iterating and scoring, especially since the field already has a similarity function stored...Why can't we just know which one to use and use it!
 - Iterating and scoring quantized vectors. By default it would be good to be able to iterate and score quantized vectors (e.g. without going through the HNSW graph).

This significantly hampers support for true exact kNN search.

This commit extends the vector value iterators to be able to return a scorer given some vector value (what this PR demonstrates). The scorer contains a copy of the originating iterator and allows for iteration and scoring the most optimized way the provided codec can give.

Users can still iterate vector values directly, read them on heap, and score any way they please.

lucene/CHANGES.txt

@@ -36,6 +36,9 @@ New Features
 * GITHUB#13288: Make HNSW and Flat storage vector formats easier to extend with new FlatVectorScorer interface. Add
   new Hnsw format for binary quantized vectors. (Ben Trent)
 
+* GITHUB#13181: Add new VectorScorer interface to vector value iterators. This allows for vector codecs to supply
+  simpler and more optimized vector scoring when iterating vector values directly. (Ben Trent)
+
 Improvements
 ---------------------
 

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

@@ -34,7 +34,9 @@
 import org.apache.lucene.index.IndexFileNames;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.KnnCollector;
+import org.apache.lucene.search.VectorScorer;
 import org.apache.lucene.store.ChecksumIndexInput;
 import org.apache.lucene.store.DataInput;
 import org.apache.lucene.store.IndexInput;
@@ -285,7 +287,8 @@ private OffHeapFloatVectorValues getOffHeapVectorValues(FieldEntry fieldEntry)
       throws IOException {
     IndexInput bytesSlice =
         vectorData.slice("vector-data", fieldEntry.vectorDataOffset, fieldEntry.vectorDataLength);
-    return new OffHeapFloatVectorValues(fieldEntry.dimension, fieldEntry.ordToDoc, bytesSlice);
+    return new OffHeapFloatVectorValues(
+        fieldEntry.dimension, fieldEntry.ordToDoc, fieldEntry.similarityFunction, bytesSlice);
   }
 
   private Bits getAcceptOrds(Bits acceptDocs, FieldEntry fieldEntry) {
@@ -365,14 +368,20 @@ static class OffHeapFloatVectorValues extends FloatVectorValues
     final int byteSize;
     int lastOrd = -1;
     final float[] value;
+    final VectorSimilarityFunction similarityFunction;
 
     int ord = -1;
     int doc = -1;
 
-    OffHeapFloatVectorValues(int dimension, int[] ordToDoc, IndexInput dataIn) {
+    OffHeapFloatVectorValues(
+        int dimension,
+        int[] ordToDoc,
+        VectorSimilarityFunction similarityFunction,
+        IndexInput dataIn) {
       this.dimension = dimension;
       this.ordToDoc = ordToDoc;
       this.dataIn = dataIn;
+      this.similarityFunction = similarityFunction;
 
       byteSize = Float.BYTES * dimension;
       value = new float[dimension];
@@ -426,7 +435,7 @@ public int advance(int target) {
 
     @Override
     public OffHeapFloatVectorValues copy() {
-      return new OffHeapFloatVectorValues(dimension, ordToDoc, dataIn.clone());
+      return new OffHeapFloatVectorValues(dimension, ordToDoc, similarityFunction, dataIn.clone());
     }
 
     @Override
@@ -439,6 +448,22 @@ public float[] vectorValue(int targetOrd) throws IOException {
       lastOrd = targetOrd;
       return value;
     }
+
+    @Override
+    public VectorScorer scorer(float[] target) {
+      OffHeapFloatVectorValues values = this.copy();
+      return new VectorScorer() {
+        @Override
+        public float score() throws IOException {
+          return values.similarityFunction.compare(values.vectorValue(), target);
+        }
+
+        @Override
+        public DocIdSetIterator iterator() {
+          return values;
+        }
+      };
+    }
   }
 
   /** Read the nearest-neighbors graph from the index input */

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

@@ -35,7 +35,9 @@
 import org.apache.lucene.index.IndexFileNames;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.KnnCollector;
+import org.apache.lucene.search.VectorScorer;
 import org.apache.lucene.store.ChecksumIndexInput;
 import org.apache.lucene.store.DataInput;
 import org.apache.lucene.store.IndexInput;
@@ -268,7 +270,11 @@ private OffHeapFloatVectorValues getOffHeapVectorValues(FieldEntry fieldEntry)
     IndexInput bytesSlice =
         vectorData.slice("vector-data", fieldEntry.vectorDataOffset, fieldEntry.vectorDataLength);
     return new OffHeapFloatVectorValues(
-        fieldEntry.dimension, fieldEntry.size(), fieldEntry.ordToDoc, bytesSlice);
+        fieldEntry.dimension,
+        fieldEntry.size(),
+        fieldEntry.ordToDoc,
+        fieldEntry.similarityFunction,
+        bytesSlice);
   }
 
   private Bits getAcceptOrds(Bits acceptDocs, FieldEntry fieldEntry) {
@@ -402,16 +408,23 @@ static class OffHeapFloatVectorValues extends FloatVectorValues
     private final IndexInput dataIn;
     private final int byteSize;
     private final float[] value;
+    private final VectorSimilarityFunction similarityFunction;
 
     private int ord = -1;
     private int doc = -1;
 
-    OffHeapFloatVectorValues(int dimension, int size, int[] ordToDoc, IndexInput dataIn) {
+    OffHeapFloatVectorValues(
+        int dimension,
+        int size,
+        int[] ordToDoc,
+        VectorSimilarityFunction similarityFunction,
+        IndexInput dataIn) {
       this.dimension = dimension;
       this.size = size;
       this.ordToDoc = ordToDoc;
       ordToDocOperator = ordToDoc == null ? IntUnaryOperator.identity() : (ord) -> ordToDoc[ord];
       this.dataIn = dataIn;
+      this.similarityFunction = similarityFunction;
       byteSize = Float.BYTES * dimension;
       value = new float[dimension];
     }
@@ -471,7 +484,8 @@ public int advance(int target) {
 
     @Override
     public OffHeapFloatVectorValues copy() {
-      return new OffHeapFloatVectorValues(dimension, size, ordToDoc, dataIn.clone());
+      return new OffHeapFloatVectorValues(
+          dimension, size, ordToDoc, similarityFunction, dataIn.clone());
     }
 
     @Override
@@ -480,6 +494,22 @@ public float[] vectorValue(int targetOrd) throws IOException {
       dataIn.readFloats(value, 0, value.length);
       return value;
     }
+
+    @Override
+    public VectorScorer scorer(float[] target) {
+      OffHeapFloatVectorValues values = this.copy();
+      return new VectorScorer() {
+        @Override
+        public float score() throws IOException {
+          return values.similarityFunction.compare(values.vectorValue(), target);
+        }
+
+        @Override
+        public DocIdSetIterator iterator() {
+          return values;
+        }
+      };
+    }
   }
 
   /** Read the nearest-neighbors graph from the index input */

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

@@ -20,6 +20,9 @@
 import java.io.IOException;
 import org.apache.lucene.codecs.lucene90.IndexedDISI;
 import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.VectorScorer;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.store.RandomAccessInput;
 import org.apache.lucene.util.Bits;
@@ -36,13 +39,19 @@ abstract class OffHeapFloatVectorValues extends FloatVectorValues
   protected final int byteSize;
   protected int lastOrd = -1;
   protected final float[] value;
+  protected final VectorSimilarityFunction vectorSimilarityFunction;
 
-  OffHeapFloatVectorValues(int dimension, int size, IndexInput slice) {
+  OffHeapFloatVectorValues(
+      int dimension,
+      int size,
+      VectorSimilarityFunction vectorSimilarityFunction,
+      IndexInput slice) {
     this.dimension = dimension;
     this.size = size;
     this.slice = slice;
     byteSize = Float.BYTES * dimension;
     value = new float[dimension];
+    this.vectorSimilarityFunction = vectorSimilarityFunction;
   }
 
   @Override
@@ -74,18 +83,24 @@ static OffHeapFloatVectorValues load(
     IndexInput bytesSlice =
         vectorData.slice("vector-data", fieldEntry.vectorDataOffset, fieldEntry.vectorDataLength);
     if (fieldEntry.docsWithFieldOffset == -1) {
-      return new DenseOffHeapVectorValues(fieldEntry.dimension, fieldEntry.size, bytesSlice);
+      return new DenseOffHeapVectorValues(
+          fieldEntry.dimension, fieldEntry.size, fieldEntry.similarityFunction, bytesSlice);
     } else {
-      return new SparseOffHeapVectorValues(fieldEntry, vectorData, bytesSlice);
+      return new SparseOffHeapVectorValues(
+          fieldEntry, vectorData, fieldEntry.similarityFunction, bytesSlice);
     }
   }
 
   static class DenseOffHeapVectorValues extends OffHeapFloatVectorValues {
 
     private int doc = -1;
 
-    public DenseOffHeapVectorValues(int dimension, int size, IndexInput slice) {
-      super(dimension, size, slice);
+    public DenseOffHeapVectorValues(
+        int dimension,
+        int size,
+        VectorSimilarityFunction vectorSimilarityFunction,
+        IndexInput slice) {
+      super(dimension, size, vectorSimilarityFunction, slice);
     }
 
     @Override
@@ -114,13 +129,29 @@ public int advance(int target) throws IOException {
 
     @Override
     public DenseOffHeapVectorValues copy() throws IOException {
-      return new DenseOffHeapVectorValues(dimension, size, slice.clone());
+      return new DenseOffHeapVectorValues(dimension, size, vectorSimilarityFunction, slice.clone());
     }
 
     @Override
     public Bits getAcceptOrds(Bits acceptDocs) {
       return acceptDocs;
     }
+
+    @Override
+    public VectorScorer scorer(float[] query) throws IOException {
+      DenseOffHeapVectorValues values = this.copy();
+      return new VectorScorer() {
+        @Override
+        public float score() throws IOException {
+          return values.vectorSimilarityFunction.compare(values.vectorValue(), query);
+        }
+
+        @Override
+        public DocIdSetIterator iterator() {
+          return values;
+        }
+      };
+    }
   }
 
   private static class SparseOffHeapVectorValues extends OffHeapFloatVectorValues {
@@ -131,10 +162,13 @@ private static class SparseOffHeapVectorValues extends OffHeapFloatVectorValues
     private final Lucene92HnswVectorsReader.FieldEntry fieldEntry;
 
     public SparseOffHeapVectorValues(
-        Lucene92HnswVectorsReader.FieldEntry fieldEntry, IndexInput dataIn, IndexInput slice)
+        Lucene92HnswVectorsReader.FieldEntry fieldEntry,
+        IndexInput dataIn,
+        VectorSimilarityFunction vectorSimilarityFunction,
+        IndexInput slice)
         throws IOException {
 
-      super(fieldEntry.dimension, fieldEntry.size, slice);
+      super(fieldEntry.dimension, fieldEntry.size, vectorSimilarityFunction, slice);
       this.fieldEntry = fieldEntry;
       final RandomAccessInput addressesData =
           dataIn.randomAccessSlice(fieldEntry.addressesOffset, fieldEntry.addressesLength);
@@ -172,8 +206,9 @@ public int advance(int target) throws IOException {
     }
 
     @Override
-    public OffHeapFloatVectorValues copy() throws IOException {
-      return new SparseOffHeapVectorValues(fieldEntry, dataIn, slice.clone());
+    public SparseOffHeapVectorValues copy() throws IOException {
+      return new SparseOffHeapVectorValues(
+          fieldEntry, dataIn, vectorSimilarityFunction, slice.clone());
     }
 
     @Override
@@ -198,12 +233,28 @@ public int length() {
         }
       };
     }
+
+    @Override
+    public VectorScorer scorer(float[] query) throws IOException {
+      SparseOffHeapVectorValues values = this.copy();
+      return new VectorScorer() {
+        @Override
+        public float score() throws IOException {
+          return values.vectorSimilarityFunction.compare(values.vectorValue(), query);
+        }
+
+        @Override
+        public DocIdSetIterator iterator() {
+          return values;
+        }
+      };
+    }
   }
 
   private static class EmptyOffHeapVectorValues extends OffHeapFloatVectorValues {
 
     public EmptyOffHeapVectorValues(int dimension) {
-      super(dimension, 0, null);
+      super(dimension, 0, VectorSimilarityFunction.COSINE, null);
     }
 
     private int doc = -1;
@@ -257,5 +308,10 @@ public int ordToDoc(int ord) {
     public Bits getAcceptOrds(Bits acceptDocs) {
       return null;
     }
+
+    @Override
+    public VectorScorer scorer(float[] query) {
+      return null;
+    }
   }
 }