Fix off-heap byte vector scoring at query time

Author: Kaival Parikh

lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/VectorScorerBenchmark.java

@@ -70,7 +70,8 @@ public class VectorScorerBenchmark {
   IndexInput in;
   KnnVectorValues vectorValues;
   byte[] vec1, vec2;
-  UpdateableRandomVectorScorer scorer;
+  UpdateableRandomVectorScorer scorer1;
+  RandomVectorScorer scorer2;
 
   @Setup(Level.Iteration)
   public void init() throws IOException {
@@ -86,11 +87,11 @@ public void init() throws IOException {
     }
     in = dir.openInput("vector.data", IOContext.DEFAULT);
     vectorValues = vectorValues(size, 2, in, DOT_PRODUCT);
-    scorer =
-        FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
-            .getRandomVectorScorerSupplier(DOT_PRODUCT, vectorValues)
-            .scorer();
-    scorer.setScoringOrdinal(0);
+
+    FlatVectorsScorer flatVectorsScorer = FlatVectorScorerUtil.getLucene99FlatVectorsScorer();
+    scorer1 = flatVectorsScorer.getRandomVectorScorerSupplier(DOT_PRODUCT, vectorValues).scorer();
+    scorer1.setScoringOrdinal(0);
+    scorer2 = flatVectorsScorer.getRandomVectorScorer(DOT_PRODUCT, vectorValues, vec1);
   }
 
   @TearDown
@@ -99,14 +100,25 @@ public void teardown() throws IOException {
   }
 
   @Benchmark
-  public float binaryDotProductDefault() throws IOException {
-    return scorer.score(1);
+  public float binaryDotProductIndexingDefault() throws IOException {
+    return scorer1.score(1);
+  }
+
+  @Benchmark
+  @Fork(jvmArgsPrepend = {"--add-modules=jdk.incubator.vector"})
+  public float binaryDotProductIndexingMemSeg() throws IOException {
+    return scorer1.score(1);
+  }
+
+  @Benchmark
+  public float binaryDotProductSearchingDefault() throws IOException {
+    return scorer2.score(1);
   }
 
   @Benchmark
   @Fork(jvmArgsPrepend = {"--add-modules=jdk.incubator.vector"})
-  public float binaryDotProductMemSeg() throws IOException {
-    return scorer.score(1);
+  public float binaryDotProductSearchingMemSeg() throws IOException {
+    return scorer2.score(1);
   }
 
   static KnnVectorValues vectorValues(

lucene/core/src/java24/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentByteVectorScorer.java

@@ -16,7 +16,10 @@
  */
 package org.apache.lucene.internal.vectorization;
 
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+
 import java.io.IOException;
+import java.lang.foreign.Arena;
 import java.lang.foreign.MemorySegment;
 import java.util.Optional;
 import org.apache.lucene.index.ByteVectorValues;
@@ -61,7 +64,7 @@ public static Optional<Lucene99MemorySegmentByteVectorScorer> create(
     super(values);
     this.input = input;
     this.vectorByteSize = values.getVectorByteLength();
-    this.query = MemorySegment.ofArray(queryVector);
+    this.query = Arena.ofAuto().allocateFrom(JAVA_BYTE, queryVector);
   }
 
   final MemorySegment getSegment(int ord) throws IOException {