Another fix; add benchmarks to cover all paths

Author: ldematte

benchmarks/src/main/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt7uBulkBenchmark.java

@@ -14,6 +14,7 @@
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.store.MMapDirectory;
+import org.apache.lucene.util.hnsw.RandomVectorScorer;
 import org.apache.lucene.util.hnsw.UpdateableRandomVectorScorer;
 import org.apache.lucene.util.quantization.QuantizedByteVectorValues;
 import org.elasticsearch.common.logging.LogConfigurator;
@@ -48,6 +49,7 @@
 import static org.elasticsearch.benchmark.vector.scorer.BenchmarkUtils.createRandomInt7VectorData;
 import static org.elasticsearch.benchmark.vector.scorer.BenchmarkUtils.getScorerFactoryOrDie;
 import static org.elasticsearch.benchmark.vector.scorer.BenchmarkUtils.luceneScoreSupplier;
+import static org.elasticsearch.benchmark.vector.scorer.BenchmarkUtils.luceneScorer;
 import static org.elasticsearch.benchmark.vector.scorer.BenchmarkUtils.readNodeCorrectionConstant;
 import static org.elasticsearch.benchmark.vector.scorer.BenchmarkUtils.supportsHeapSegments;
 import static org.elasticsearch.benchmark.vector.scorer.BenchmarkUtils.vectorValues;
@@ -100,6 +102,9 @@ public class VectorScorerInt7uBulkBenchmark {
     UpdateableRandomVectorScorer luceneDotScorer;
     UpdateableRandomVectorScorer nativeDotScorer;
 
+    RandomVectorScorer luceneDotScorerQuery;
+    RandomVectorScorer nativeDotScorerQuery;
+
     @Setup(Level.Trial)
     public void setup() throws IOException {
         factory = getScorerFactoryOrDie();
@@ -127,6 +132,17 @@ public void setup() throws IOException {
             .orElseThrow()
             .scorer();
         nativeDotScorer.setScoringOrdinal(targetOrd);
+
+        if (supportsHeapSegments()) {
+            // setup for getInt7SQVectorScorer / query vector scoring
+            float[] queryVec = new float[dims];
+            for (int i = 0; i < dims; i++) {
+                queryVec[i] = random.nextFloat();
+            }
+            luceneDotScorerQuery = luceneScorer(dotProductValues, VectorSimilarityFunction.DOT_PRODUCT, queryVec);
+            nativeDotScorerQuery = factory.getInt7SQVectorScorer(VectorSimilarityFunction.DOT_PRODUCT, dotProductValues, queryVec)
+                .orElseThrow();
+        }
     }
 
     @TearDown
@@ -151,6 +167,14 @@ public float[] dotProductLuceneMultipleRandom() throws IOException {
         return scores;
     }
 
+    @Benchmark
+    public float[] dotProductLuceneQueryMultipleRandom() throws IOException {
+        for (int v = 0; v < numVectorsToScore; v++) {
+            scores[v] = luceneDotScorerQuery.score(ordinals[v]);
+        }
+        return scores;
+    }
+
     @Benchmark
     public float[] dotProductNativeMultipleSequential() throws IOException {
         for (int v = 0; v < numVectorsToScore; v++) {
@@ -167,6 +191,14 @@ public float[] dotProductNativeMultipleRandom() throws IOException {
         return scores;
     }
 
+    @Benchmark
+    public float[] dotProductNativeQueryMultipleRandom() throws IOException {
+        for (int v = 0; v < numVectorsToScore; v++) {
+            scores[v] = nativeDotScorerQuery.score(ordinals[v]);
+        }
+        return scores;
+    }
+
     @Benchmark
     public float[] dotProductNativeMultipleSequentialBulk() throws IOException {
         nativeDotScorer.bulkScore(ids, scores, ordinals.length);
@@ -179,6 +211,12 @@ public float[] dotProductNativeMultipleRandomBulk() throws IOException {
         return scores;
     }
 
+    @Benchmark
+    public float[] dotProductNativeQueryMultipleRandomBulk() throws IOException {
+        nativeDotScorerQuery.bulkScore(ordinals, scores, ordinals.length);
+        return scores;
+    }
+
     @Benchmark
     public float[] dotProductScalarMultipleSequential() throws IOException {
         var queryVector = dotProductValues.vectorValue(targetOrd);

benchmarks/src/test/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt7uBulkBenchmarkTests.java

@@ -18,6 +18,8 @@
 
 import java.util.Arrays;
 
+import static org.elasticsearch.benchmark.vector.scorer.BenchmarkUtils.supportsHeapSegments;
+
 public class VectorScorerInt7uBulkBenchmarkTests extends ESTestCase {
 
     final float delta = 1e-3f;
@@ -61,6 +63,11 @@ public void testDotProductRandom() throws Exception {
                 assertArrayEquals(expected, bench.dotProductLuceneMultipleRandom(), delta);
                 assertArrayEquals(expected, bench.dotProductNativeMultipleRandom(), delta);
                 assertArrayEquals(expected, bench.dotProductNativeMultipleRandomBulk(), delta);
+                if (supportsHeapSegments()) {
+                    assertArrayEquals(expected, bench.dotProductLuceneQueryMultipleRandom(), delta);
+                    assertArrayEquals(expected, bench.dotProductNativeQueryMultipleRandom(), delta);
+                    assertArrayEquals(expected, bench.dotProductNativeQueryMultipleRandomBulk(), delta);
+                }
             } finally {
                 bench.teardown();
             }

libs/native/src/main/java/org/elasticsearch/nativeaccess/jdk/JdkVectorLibrary.java

@@ -61,7 +61,7 @@ public final class JdkVectorLibrary implements VectorLibrary {
                         LinkerHelperUtil.critical()
                     );
                     dot7uBulkWithOffsets$mh = downcallHandle(
-                        "dot7u_bulk_offsets_2",
+                        "vec_dot7u_bulk_offsets_2",
                         FunctionDescriptor.ofVoid(ADDRESS, ADDRESS, JAVA_INT, ADDRESS, JAVA_INT, ADDRESS),
                         LinkerHelperUtil.critical()
                     );
@@ -97,7 +97,7 @@ public final class JdkVectorLibrary implements VectorLibrary {
                         LinkerHelperUtil.critical()
                     );
                     dot7uBulkWithOffsets$mh = downcallHandle(
-                        "dot7u_bulk_offsets",
+                        "vec_dot7u_bulk_offsets",
                         FunctionDescriptor.ofVoid(ADDRESS, ADDRESS, JAVA_INT, ADDRESS, JAVA_INT, ADDRESS),
                         LinkerHelperUtil.critical()
                     );