Add DirectIO bulk rescoring

Author: benwtrent

server/src/main/java/org/elasticsearch/index/codec/vectors/BulkScorableFloatVectorValues.java

@@ -19,11 +19,11 @@ public interface BulkScorableFloatVectorValues extends BulkScorableVectorValues
      * Returns a {@link BulkVectorScorer} that can score against the provided {@code target} vector.
      * It will score to the fastest speed possible, potentially sacrificing some fidelity.
      */
-    BulkVectorScorer scorer(float[] target) throws IOException;
+    BulkVectorScorer bulkScorer(float[] target) throws IOException;
 
     /**
      * Returns a {@link BulkVectorScorer} that can rescore against the provided {@code target} vector.
      * It will score to the highest fidelity possible, potentially sacrificing some speed.
      */
-    BulkVectorScorer rescorer(float[] target) throws IOException;
+    BulkVectorScorer bulkRescorer(float[] target) throws IOException;
 }

server/src/main/java/org/elasticsearch/index/codec/vectors/BulkScorableVectorValues.java

@@ -30,9 +30,8 @@ interface BulkVectorScorer extends VectorScorer {
         interface Bulk {
             /**
              * Scores up to {@code nextCount} docs in the provided {@code buffer}.
-             * Returns the maxScore of docs scored.
              */
-            float nextDocsAndScores(int nextCount, Bits liveDocs, DocAndFloatFeatureBuffer buffer) throws IOException;
+            void nextDocsAndScores(int nextCount, Bits liveDocs, DocAndFloatFeatureBuffer buffer) throws IOException;
         }
     }
 }

server/src/main/java/org/elasticsearch/index/codec/vectors/es818/DirectIOLucene99FlatVectorsFormat.java

@@ -22,19 +22,35 @@
 import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
+import org.apache.lucene.codecs.lucene95.HasIndexSlice;
 import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsReader;
 import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsWriter;
+import org.apache.lucene.index.ByteVectorValues;
+import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.index.KnnVectorValues;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.ConjunctionUtils;
+import org.apache.lucene.search.DocAndFloatFeatureBuffer;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.VectorScorer;
 import org.apache.lucene.store.FlushInfo;
 import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.store.MergeInfo;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.hnsw.RandomVectorScorer;
 import org.elasticsearch.common.util.set.Sets;
 import org.elasticsearch.index.codec.vectors.AbstractFlatVectorsFormat;
+import org.elasticsearch.index.codec.vectors.BulkScorableFloatVectorValues;
+import org.elasticsearch.index.codec.vectors.BulkScorableVectorValues;
 import org.elasticsearch.index.codec.vectors.MergeReaderWrapper;
 import org.elasticsearch.index.store.FsDirectoryFactory;
 
 import java.io.IOException;
+import java.util.List;
 import java.util.Set;
 
 /**
@@ -86,8 +102,13 @@ public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException
             );
             // Use mmap for merges and direct I/O for searches.
             // TODO: Open the mmap file with sequential access instead of random (current behavior).
+            // TODO: maybe we should force completely RANDOM access always for inner reader formats (outside of merges)?
             return new MergeReaderWrapper(
-                new Lucene99FlatVectorsReader(directIOState, vectorsScorer),
+                new Lucene99FlatBulkScoringVectorsReader(
+                    directIOState,
+                    new Lucene99FlatVectorsReader(directIOState, vectorsScorer),
+                    vectorsScorer
+                ),
                 new Lucene99FlatVectorsReader(state, vectorsScorer)
             );
         } else {
@@ -129,4 +150,204 @@ public IOContext withHints(FileOpenHint... hints) {
             return new DirectIOContext(Set.of(hints));
         }
     }
+
+    static class Lucene99FlatBulkScoringVectorsReader extends FlatVectorsReader {
+        private final Lucene99FlatVectorsReader inner;
+        private final SegmentReadState state;
+
+        Lucene99FlatBulkScoringVectorsReader(SegmentReadState state, Lucene99FlatVectorsReader inner, FlatVectorsScorer scorer) {
+            super(scorer);
+            this.inner = inner;
+            this.state = state;
+        }
+
+        @Override
+        public void close() throws IOException {
+            inner.close();
+        }
+
+        @Override
+        public RandomVectorScorer getRandomVectorScorer(String field, float[] target) throws IOException {
+            return inner.getRandomVectorScorer(field, target);
+        }
+
+        @Override
+        public RandomVectorScorer getRandomVectorScorer(String field, byte[] target) throws IOException {
+            return inner.getRandomVectorScorer(field, target);
+        }
+
+        @Override
+        public void checkIntegrity() throws IOException {
+            inner.checkIntegrity();
+        }
+
+        @Override
+        public FloatVectorValues getFloatVectorValues(String field) throws IOException {
+            FloatVectorValues vectorValues = inner.getFloatVectorValues(field);
+            if (vectorValues == null || vectorValues.size() == 0) {
+                return null;
+            }
+            FieldInfo info = state.fieldInfos.fieldInfo(field);
+            return new RescorerOffHeapVectorValues(vectorValues, info.getVectorSimilarityFunction(), vectorScorer);
+        }
+
+        @Override
+        public ByteVectorValues getByteVectorValues(String field) throws IOException {
+            return inner.getByteVectorValues(field);
+        }
+
+        @Override
+        public long ramBytesUsed() {
+            return inner.ramBytesUsed();
+        }
+
+        static class RescorerOffHeapVectorValues extends FloatVectorValues implements BulkScorableFloatVectorValues {
+            VectorSimilarityFunction similarityFunction;
+            FloatVectorValues inner;
+            IndexInput inputSlice;
+            FlatVectorsScorer scorer;
+
+            RescorerOffHeapVectorValues(FloatVectorValues inner, VectorSimilarityFunction similarityFunction, FlatVectorsScorer scorer) {
+                this.inner = inner;
+                if (inner instanceof HasIndexSlice slice) {
+                    this.inputSlice = slice.getSlice();
+                } else {
+                    this.inputSlice = null;
+                }
+                this.similarityFunction = similarityFunction;
+                this.scorer = scorer;
+            }
+
+            @Override
+            public float[] vectorValue(int ord) throws IOException {
+                return inner.vectorValue(ord);
+            }
+
+            @Override
+            public int dimension() {
+                return inner.dimension();
+            }
+
+            @Override
+            public int size() {
+                return inner.size();
+            }
+
+            @Override
+            public RescorerOffHeapVectorValues copy() throws IOException {
+                return new RescorerOffHeapVectorValues(inner.copy(), similarityFunction, scorer);
+            }
+
+            @Override
+            public BulkVectorScorer bulkRescorer(float[] target) throws IOException {
+                return bulkScorer(target);
+            }
+
+            @Override
+            public BulkVectorScorer bulkScorer(float[] target) throws IOException {
+                DocIndexIterator indexIterator = inner.iterator();
+                RandomVectorScorer randomScorer = scorer.getRandomVectorScorer(similarityFunction, inner, target);
+                return new PreFetchingFloatBulkScorer(randomScorer, indexIterator, inputSlice, dimension() * Float.BYTES);
+            }
+
+            @Override
+            public VectorScorer scorer(float[] target) throws IOException {
+                return inner.scorer(target);
+            }
+        }
+
+        private record PreFetchingFloatBulkScorer(
+            RandomVectorScorer inner,
+            KnnVectorValues.DocIndexIterator indexIterator,
+            IndexInput inputSlice,
+            int byteSize
+        ) implements BulkScorableVectorValues.BulkVectorScorer {
+
+            @Override
+            public float score() throws IOException {
+                return inner.score(indexIterator.index());
+            }
+
+            @Override
+            public DocIdSetIterator iterator() {
+                return indexIterator;
+            }
+
+            @Override
+            public Bulk bulk(DocIdSetIterator matchingDocs) throws IOException {
+                DocIdSetIterator conjunctionScorer = matchingDocs == null
+                    ? indexIterator
+                    : ConjunctionUtils.intersectIterators(List.of(matchingDocs, indexIterator));
+                if (conjunctionScorer.docID() == -1) {
+                    conjunctionScorer.nextDoc();
+                }
+                return new FloatBulkScorer(inner, inputSlice, byteSize, 32, indexIterator, conjunctionScorer);
+            }
+        }
+
+        private static class FloatBulkScorer implements BulkScorableVectorValues.BulkVectorScorer.Bulk {
+            private final KnnVectorValues.DocIndexIterator indexIterator;
+            private final DocIdSetIterator matchingDocs;
+            private final RandomVectorScorer inner;
+            private final int bulkSize;
+            private final IndexInput inputSlice;
+            private final int byteSize;
+            private final int[] docBuffer;
+            private final float[] scoreBuffer;
+
+            FloatBulkScorer(
+                RandomVectorScorer fvv,
+                IndexInput inputSlice,
+                int byteSize,
+                int bulkSize,
+                KnnVectorValues.DocIndexIterator iterator,
+                DocIdSetIterator matchingDocs
+            ) {
+                this.indexIterator = iterator;
+                this.matchingDocs = matchingDocs;
+                this.inner = fvv;
+                this.bulkSize = bulkSize;
+                this.inputSlice = inputSlice;
+                this.docBuffer = new int[bulkSize];
+                this.scoreBuffer = new float[bulkSize];
+                this.byteSize = byteSize;
+            }
+
+            @Override
+            public void nextDocsAndScores(int nextCount, Bits liveDocs, DocAndFloatFeatureBuffer buffer) throws IOException {
+                buffer.growNoCopy(nextCount);
+                int size = 0;
+                for (int doc = matchingDocs.docID(); doc != DocIdSetIterator.NO_MORE_DOCS && size < nextCount; doc = matchingDocs
+                    .nextDoc()) {
+                    if (liveDocs == null || liveDocs.get(doc)) {
+                        buffer.docs[size++] = indexIterator.index();
+                    }
+                }
+                int loopBound = size - (size % bulkSize);
+                int i = 0;
+                for (; i < loopBound; i += bulkSize) {
+                    for (int j = 0; j < bulkSize; j++) {
+                        long ord = buffer.docs[i + j];
+                        inputSlice.prefetch(ord * byteSize, byteSize);
+                    }
+                    System.arraycopy(buffer.docs, i, docBuffer, 0, bulkSize);
+                    inner.bulkScore(docBuffer, scoreBuffer, bulkSize);
+                    System.arraycopy(scoreBuffer, 0, buffer.features, i, bulkSize);
+                }
+                int countLeft = size - i;
+                for (int j = i; j < size; j++) {
+                    long ord = buffer.docs[j];
+                    inputSlice.prefetch(ord * byteSize, byteSize);
+                }
+                System.arraycopy(buffer.docs, i, docBuffer, 0, countLeft);
+                inner.bulkScore(docBuffer, scoreBuffer, countLeft);
+                System.arraycopy(scoreBuffer, 0, buffer.features, i, countLeft);
+                buffer.size = size;
+                // fix the docIds in buffer
+                for (int j = 0; j < size; j++) {
+                    buffer.docs[j] = inner.ordToDoc(buffer.docs[j]);
+                }
+            }
+        }
+    }
 }

server/src/main/java/org/elasticsearch/search/vectors/RescoreKnnVectorQuery.java

@@ -317,7 +317,7 @@ private void rescoreBulk(
             List<ScoreDoc> queue,
             DocIdSetIterator filterIterator
         ) throws IOException {
-            BulkScorableVectorValues.BulkVectorScorer vectorReScorer = rescorableVectorValues.rescorer(floatTarget);
+            BulkScorableVectorValues.BulkVectorScorer vectorReScorer = rescorableVectorValues.bulkRescorer(floatTarget);
             var iterator = vectorReScorer.iterator();
             BulkScorableVectorValues.BulkVectorScorer.Bulk bulkScorer = vectorReScorer.bulk(filterIterator);
             DocAndFloatFeatureBuffer buffer = new DocAndFloatFeatureBuffer();