Introduce NativeEngines990KnnVectorsScorer to decouple native SIMD scoring selection from FaissMemoryOptimizedSearcher

Previously, FaissMemoryOptimizedSearcher was responsible for detecting
MMap-backed vectors and choosing between NativeRandomVectorScorer and
the default Java scorer. This logic is now extracted into a dedicated
FlatVectorsScorer decorator that sits in the codec scoring chain,
making native SIMD acceleration available to all scoring paths
(including HNSW graph traversal with prefetch support).

By placing NativeEngines990KnnVectorsScorer inside
PrefetchableFlatVectorScorer in the decorator chain, the native scorer
now benefits from prefetch-enabled bulk scoring during HNSW graph
traversal. Previously, NativeRandomVectorScorer was instantiated
directly in FaissMemoryOptimizedSearcher, bypassing the prefetch layer
entirely. With this change, PrefetchableFlatVectorScorer can wrap the
NativeRandomVectorScorer returned by NativeEngines990KnnVectorsScorer,
issuing prefetch hints for memory-mapped vector data before native SIMD
scoring, reducing memory access latency during graph neighbor evaluation.

Key changes:
- Add NativeEngines990KnnVectorsScorer which wraps a FlatVectorsScorer
  delegate and transparently returns NativeRandomVectorScorer when the
  bottom-level FloatVectorValues implements MMapVectorValues and the
  similarity function is EUCLIDEAN or MAXIMUM_INNER_PRODUCT
- Simplify FaissMemoryOptimizedSearcher by removing the native scoring
  branch and determineNativeFunctionType(); it now unconditionally
  delegates to the FlatVectorsScorer
- Refactor NativeRandomVectorScorer to extend AbstractRandomVectorScorer,
  removing manual maxOrd/ordToDoc/getAcceptOrds overrides, enabling it
  to be wrapped by PrefetchableFlatVectorScorer for prefetch-enabled
  bulk scoring during HNSW graph traversal
- Wire the new scorer into NativeEngines990KnnVectorsFormat between the
  Lucene99 flat vectors scorer and PrefetchableFlatVectorScorer
- Add unit tests covering all routing branches

Signed-off-by: Vijayan Balasubramanian <balasvij@amazon.com>

Author: VijayanB

CHANGELOG.md

@@ -32,3 +32,4 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 * Add Prefetch functionality to prefetch vectors during ANN Search for MemoryOptimizedSearch. [#3173](https://github.com/opensearch-project/k-NN/pull/3173)
 * Optimize ByteVectorIdsExactKNNIterator by moving array conversion to constructor [#3171](https://github.com/opensearch-project/k-NN/pull/3171)
 * Add VectorScorers for BinaryDocValues and nested best child scoring [#3179](https://github.com/opensearch-project/k-NN/pull/3179)
+* Introduce NativeEngines990KnnVectorsScorer to decouple native SIMD scoring selection from FaissMemoryOptimizedSearcher [#3184](https://github.com/opensearch-project/k-NN/pull/3184)

src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsFormat.java

@@ -22,6 +22,7 @@
 import org.apache.lucene.index.SegmentWriteState;
 import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.codec.nativeindex.NativeIndexBuildStrategyFactory;
+import org.opensearch.knn.index.codec.scorer.NativeEngines990KnnVectorsScorer;
 import org.opensearch.knn.index.codec.scorer.PrefetchableFlatVectorScorer;
 import org.opensearch.knn.index.engine.KNNEngine;
 
@@ -35,7 +36,7 @@
 public class NativeEngines990KnnVectorsFormat extends KnnVectorsFormat {
     /** The format for storing, reading, merging vectors on disk */
     private static final FlatVectorsFormat flatVectorsFormat = new Lucene99FlatVectorsFormat(
-        new PrefetchableFlatVectorScorer(FlatVectorScorerUtil.getLucene99FlatVectorsScorer())
+        new PrefetchableFlatVectorScorer(new NativeEngines990KnnVectorsScorer(FlatVectorScorerUtil.getLucene99FlatVectorsScorer()))
     );
     private static final String FORMAT_NAME = "NativeEngines990KnnVectorsFormat";
     private final int approximateThreshold;

src/main/java/org/opensearch/knn/index/codec/scorer/NativeEngines990KnnVectorsScorer.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.index.codec.scorer;
+
+import lombok.RequiredArgsConstructor;
+import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
+import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.index.KnnVectorValues;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.util.hnsw.RandomVectorScorer;
+import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier;
+import org.opensearch.knn.jni.SimdVectorComputeService;
+import org.opensearch.knn.memoryoptsearch.faiss.MMapVectorValues;
+import org.opensearch.knn.memoryoptsearch.faiss.NativeRandomVectorScorer;
+import org.opensearch.knn.memoryoptsearch.faiss.WrappedFloatVectorValues;
+
+import java.io.IOException;
+
+/**
+ * A {@link FlatVectorsScorer} that transparently selects between native SIMD-optimized scoring
+ * and pure-Java scoring based on whether the underlying vector values are memory-mapped.
+ *
+ * <p>When the bottom-level {@link FloatVectorValues} implements {@link MMapVectorValues},
+ * this scorer returns a {@link NativeRandomVectorScorer} for hardware-accelerated computation.
+ * Otherwise, it delegates to the wrapped {@link FlatVectorsScorer}.
+ */
+@RequiredArgsConstructor
+public class NativeEngines990KnnVectorsScorer implements FlatVectorsScorer {
+    private final FlatVectorsScorer delegate;
+
+    @Override
+    public RandomVectorScorer getRandomVectorScorer(
+        VectorSimilarityFunction similarityFunction,
+        KnnVectorValues vectorValues,
+        float[] target
+    ) throws IOException {
+        final SimdVectorComputeService.SimilarityFunctionType nativeType = getNativeFunctionType(similarityFunction);
+        if (nativeType != null) {
+            final FloatVectorValues bottomValues = WrappedFloatVectorValues.getBottomFloatVectorValues(vectorValues);
+            if (bottomValues instanceof MMapVectorValues mmapValues) {
+                return new NativeRandomVectorScorer(target, vectorValues, mmapValues, nativeType);
+            }
+        }
+        return delegate.getRandomVectorScorer(similarityFunction, vectorValues, target);
+    }
+
+    @Override
+    public RandomVectorScorer getRandomVectorScorer(
+        VectorSimilarityFunction similarityFunction,
+        KnnVectorValues vectorValues,
+        byte[] target
+    ) throws IOException {
+        return delegate.getRandomVectorScorer(similarityFunction, vectorValues, target);
+    }
+
+    @Override
+    public RandomVectorScorerSupplier getRandomVectorScorerSupplier(
+        VectorSimilarityFunction similarityFunction,
+        KnnVectorValues vectorValues
+    ) throws IOException {
+        return delegate.getRandomVectorScorerSupplier(similarityFunction, vectorValues);
+    }
+
+    private static SimdVectorComputeService.SimilarityFunctionType getNativeFunctionType(
+        final VectorSimilarityFunction similarityFunction
+    ) {
+        return switch (similarityFunction) {
+            case MAXIMUM_INNER_PRODUCT -> SimdVectorComputeService.SimilarityFunctionType.FP16_MAXIMUM_INNER_PRODUCT;
+            case EUCLIDEAN -> SimdVectorComputeService.SimilarityFunctionType.FP16_L2;
+            default -> null;
+        };
+    }
+}

src/main/java/org/opensearch/knn/memoryoptsearch/faiss/FaissMemoryOptimizedSearcher.java

@@ -9,7 +9,6 @@
 import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
 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.VectorEncoding;
 import org.apache.lucene.index.VectorSimilarityFunction;
@@ -29,7 +28,6 @@
 import org.opensearch.knn.index.KNNVectorSimilarityFunction;
 import org.opensearch.knn.index.SpaceType;
 import org.opensearch.knn.index.engine.qframe.QuantizationConfig;
-import org.opensearch.knn.jni.SimdVectorComputeService;
 import org.opensearch.knn.memoryoptsearch.VectorSearcher;
 import org.opensearch.knn.memoryoptsearch.faiss.cagra.FaissCagraHNSW;
 
@@ -58,7 +56,6 @@ public WarmupInitializationException(String message) {
     private final VectorSimilarityFunction vectorSimilarityFunction;
     private final long fileSize;
     private boolean isAdc;
-    private SimdVectorComputeService.SimilarityFunctionType nativeSimilarityFunctionType;
 
     public FaissMemoryOptimizedSearcher(final IndexInput indexInput, final FieldInfo fieldInfo, final FlatVectorsScorer flatVectorsScorer)
         throws IOException {
@@ -90,7 +87,6 @@ public FaissMemoryOptimizedSearcher(final IndexInput indexInput, final FieldInfo
         );
 
         this.hnsw = extractFaissHnsw(faissIndex);
-        this.nativeSimilarityFunctionType = determineNativeFunctionType();
     }
 
     private static FaissHNSW extractFaissHnsw(final FaissIndex faissIndex) {
@@ -106,35 +102,13 @@ public void search(float[] target, KnnCollector knnCollector, AcceptDocs acceptD
         final KnnVectorValues knnVectorValues = isAdc
             ? faissIndex.getByteValues(getSlicedIndexInput())
             : faissIndex.getFloatValues(getSlicedIndexInput());
-        final FloatVectorValues bottomKnnVectorValues = WrappedFloatVectorValues.getBottomFloatVectorValues(knnVectorValues);
-        final boolean useNativeScoring = bottomKnnVectorValues instanceof MMapVectorValues;
-        final IOSupplier<RandomVectorScorer> scorerSupplier;
-
-        if (useNativeScoring) {
-            // We can use native scoring.
-            scorerSupplier = () -> new NativeRandomVectorScorer(
-                target,
-                knnVectorValues,
-                (MMapVectorValues) bottomKnnVectorValues,
-                nativeSimilarityFunctionType
-            );
-        } else {
-            // Falling back to default scoring using pure Java.
-            scorerSupplier = () -> flatVectorsScorer.getRandomVectorScorer(vectorSimilarityFunction, knnVectorValues, target);
-        }
 
-        search(VectorEncoding.FLOAT32, scorerSupplier, knnCollector, acceptDocs);
-    }
-
-    private SimdVectorComputeService.SimilarityFunctionType determineNativeFunctionType() {
-        if (vectorSimilarityFunction == VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT) {
-            return SimdVectorComputeService.SimilarityFunctionType.FP16_MAXIMUM_INNER_PRODUCT;
-        } else if (vectorSimilarityFunction == VectorSimilarityFunction.EUCLIDEAN) {
-            return SimdVectorComputeService.SimilarityFunctionType.FP16_L2;
-        }
-
-        // At the moment, we only support FP16, it's fine to return null.
-        return null;
+        search(
+            VectorEncoding.FLOAT32,
+            () -> flatVectorsScorer.getRandomVectorScorer(vectorSimilarityFunction, knnVectorValues, target),
+            knnCollector,
+            acceptDocs
+        );
     }
 
     @Override

src/main/java/org/opensearch/knn/memoryoptsearch/faiss/NativeRandomVectorScorer.java

@@ -5,9 +5,7 @@
 
 package org.opensearch.knn.memoryoptsearch.faiss;
 
-import lombok.NonNull;
 import org.apache.lucene.index.KnnVectorValues;
-import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.hnsw.RandomVectorScorer;
 import org.opensearch.knn.jni.SimdVectorComputeService;
 
@@ -21,21 +19,12 @@
  * memory-mapped vector chunks, and delegates all similarity scoring operations to the
  * {@link SimdVectorComputeService}. The underlying native library is expected to
  * leverage SIMD instructions (e.g., AVX, AVX512, or NEON) to accelerate computations.
+ * <p>
+ * Extends {@link AbstractRandomVectorScorer} so that it can be wrapped by
+ * {@link org.opensearch.knn.index.codec.scorer.PrefetchableFlatVectorScorer} for
+ * prefetch-enabled bulk scoring during HNSW graph traversal.
  */
-public class NativeRandomVectorScorer implements RandomVectorScorer {
-
-    // Backing {@link KnnVectorValues} used for document–vector association.
-    @NonNull
-    private final KnnVectorValues knnVectorValues;
-
-    // Array of address–size pairs describing memory-mapped vector chunks.
-    private long[] addressAndSize;
-
-    // Maximum vector id available for scoring.
-    private int maxOrd;
-
-    // Index value of the native similarity function type.
-    private int nativeFunctionTypeOrd;
+public class NativeRandomVectorScorer extends RandomVectorScorer.AbstractRandomVectorScorer {
 
     /**
      * Constructs a native-backed scorer for computing similarity between the given query
@@ -52,11 +41,8 @@ public NativeRandomVectorScorer(
         final MMapVectorValues mmapVectorValues,
         final SimdVectorComputeService.SimilarityFunctionType similarityFunctionType
     ) {
-        this.knnVectorValues = knnVectorValues;
-        this.addressAndSize = mmapVectorValues.getAddressAndSize();
-        this.maxOrd = knnVectorValues.size();
-        this.nativeFunctionTypeOrd = similarityFunctionType.ordinal();
-        SimdVectorComputeService.saveSearchContext(query, addressAndSize, nativeFunctionTypeOrd);
+        super(knnVectorValues);
+        SimdVectorComputeService.saveSearchContext(query, mmapVectorValues.getAddressAndSize(), similarityFunctionType.ordinal());
     }
 
     /**
@@ -87,36 +73,4 @@ public float bulkScore(final int[] internalVectorIds, final float[] scores, fina
     public float score(final int internalVectorId) throws IOException {
         return SimdVectorComputeService.scoreSimilarity(internalVectorId);
     }
-
-    /**
-     * Returns the maximum vector id for scoring.
-     *
-     * @return the maximum vector id
-     */
-    @Override
-    public int maxOrd() {
-        return maxOrd;
-    }
-
-    /**
-     * Maps an internal vector ordinal to its corresponding document ID.
-     *
-     * @param ord the internal vector id
-     * @return the document ID associated with the given vector id
-     */
-    @Override
-    public int ordToDoc(int ord) {
-        return knnVectorValues.ordToDoc(ord);
-    }
-
-    /**
-     * Returns a filtered {@link Bits} view representing accepted documents.
-     *
-     * @param acceptDocs the bit set of accepted documents
-     * @return a {@link Bits} object describing acceptable vector ids for scoring
-     */
-    @Override
-    public Bits getAcceptOrds(Bits acceptDocs) {
-        return knnVectorValues.getAcceptOrds(acceptDocs);
-    }
 }