Feature/scalar quantized off heap scoring

lucene/core/src/java/org/apache/lucene/codecs/hnsw/FlatVectorScorerUtil.java

@@ -37,4 +37,12 @@ private FlatVectorScorerUtil() {}
   public static FlatVectorsScorer getLucene99FlatVectorsScorer() {
     return IMPL.getLucene99FlatVectorsScorer();
   }
+
+  /**
+   * Returns a FlatVectorsScorer that supports the Lucene99 format. Scorers retrieved through this
+   * method may be optimized on certain platforms. Otherwise, a DefaultFlatVectorScorer is returned.
+   */
+  public static FlatVectorsScorer getLucene99ScalarQuantizedVectorsScorer() {
+    return IMPL.getLucene99ScalarQuantizedVectorsScorer();
+  }
 }

lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99ScalarQuantizedVectorsFormat.java

@@ -18,10 +18,10 @@
 package org.apache.lucene.codecs.lucene99;
 
 import java.io.IOException;
-import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer;
 import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil;
 import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
 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.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
@@ -68,7 +68,7 @@ public class Lucene99ScalarQuantizedVectorsFormat extends FlatVectorsFormat {
 
   final byte bits;
   final boolean compress;
-  final Lucene99ScalarQuantizedVectorScorer flatVectorScorer;
+  final FlatVectorsScorer flatVectorScorer;
 
   /** Constructs a format using default graph construction parameters */
   public Lucene99ScalarQuantizedVectorsFormat() {
@@ -109,8 +109,7 @@ public Lucene99ScalarQuantizedVectorsFormat(
     this.bits = (byte) bits;
     this.confidenceInterval = confidenceInterval;
     this.compress = compress;
-    this.flatVectorScorer =
-        new Lucene99ScalarQuantizedVectorScorer(DefaultFlatVectorScorer.INSTANCE);
+    this.flatVectorScorer = FlatVectorScorerUtil.getLucene99ScalarQuantizedVectorsScorer();
   }
 
   public static float calculateDefaultConfidenceInterval(int vectorDimension) {

lucene/core/src/java/org/apache/lucene/internal/vectorization/DefaultVectorizationProvider.java

@@ -19,6 +19,7 @@
 
 import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer;
 import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
+import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorScorer;
 import org.apache.lucene.store.IndexInput;
 
 /** Default provider returning scalar implementations. */
@@ -40,6 +41,11 @@ public FlatVectorsScorer getLucene99FlatVectorsScorer() {
     return DefaultFlatVectorScorer.INSTANCE;
   }
 
+  @Override
+  public FlatVectorsScorer getLucene99ScalarQuantizedVectorsScorer() {
+    return new Lucene99ScalarQuantizedVectorScorer(getLucene99FlatVectorsScorer());
+  }
+
   @Override
   public PostingDecodingUtil newPostingDecodingUtil(IndexInput input) {
     return new DefaultPostingDecodingUtil(input);

lucene/core/src/java/org/apache/lucene/internal/vectorization/VectorizationProvider.java

@@ -97,6 +97,9 @@ public static VectorizationProvider getInstance() {
   /** Returns a FlatVectorsScorer that supports the Lucene99 format. */
   public abstract FlatVectorsScorer getLucene99FlatVectorsScorer();
 
+  /** Returns a FlatVectorsScorer that supports scalar quantized vectors in the Lucene99 format. */
+  public abstract FlatVectorsScorer getLucene99ScalarQuantizedVectorsScorer();
+
   /** Create a new {@link PostingDecodingUtil} for the given {@link IndexInput}. */
   public abstract PostingDecodingUtil newPostingDecodingUtil(IndexInput input) throws IOException;
 

lucene/core/src/java21/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentScalarQuantizedFlatVectorsScorer.java

@@ -0,0 +1,123 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.internal.vectorization;
+
+import static org.apache.lucene.codecs.hnsw.ScalarQuantizedVectorScorer.quantizeQuery;
+
+import java.io.IOException;
+import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.util.hnsw.RandomAccessVectorValues;
+import org.apache.lucene.util.hnsw.RandomVectorScorer;
+import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier;
+import org.apache.lucene.util.quantization.RandomAccessQuantizedByteVectorValues;
+
+public class Lucene99MemorySegmentScalarQuantizedFlatVectorsScorer implements FlatVectorsScorer {
+
+  private final FlatVectorsScorer delegate;
+
+  public Lucene99MemorySegmentScalarQuantizedFlatVectorsScorer(FlatVectorsScorer delegate) {
+    this.delegate = delegate;
+  }
+
+  @Override
+  public RandomVectorScorerSupplier getRandomVectorScorerSupplier(
+      VectorSimilarityFunction similarityType, RandomAccessVectorValues vectorValues)
+      throws IOException {
+    if (vectorValues instanceof RandomAccessQuantizedByteVectorValues quantizedByteVectorValues) {
+      // Unoptimized edge case, we don't optimize compressed 4-bit quantization with Euclidean
+      // similarity
+      // So, we delegate to the default scorer
+      if (quantizedByteVectorValues.getScalarQuantizer().getBits() == 4
+          && similarityType == VectorSimilarityFunction.EUCLIDEAN
+          // Indicates that the vector is compressed as the byte length is not equal to the
+          // dimension count
+          && vectorValues.getVectorByteLength() != vectorValues.dimension()) {
+        return delegate.getRandomVectorScorerSupplier(similarityType, vectorValues);
+      }
+      var scalarQuantizer = quantizedByteVectorValues.getScalarQuantizer();
+      var scalarScorerSupplier =
+          Lucene99MemorySegmentScalarQuantizedVectorScorerSupplier.create(
+              similarityType,
+              scalarQuantizer.getBits(),
+              scalarQuantizer.getConstantMultiplier(),
+              quantizedByteVectorValues);
+      if (scalarScorerSupplier.isPresent()) {
+        return scalarScorerSupplier.get();
+      }
+    }
+    return delegate.getRandomVectorScorerSupplier(similarityType, vectorValues);
+  }
+
+  @Override
+  public RandomVectorScorer getRandomVectorScorer(
+      VectorSimilarityFunction similarityType,
+      RandomAccessVectorValues vectorValues,
+      float[] queryVector)
+      throws IOException {
+    if (vectorValues instanceof RandomAccessQuantizedByteVectorValues quantizedByteVectorValues) {
+      // Unoptimized edge case, we don't optimize compressed 4-bit quantization with Euclidean
+      // similarity
+      // So, we delegate to the default scorer
+      if (quantizedByteVectorValues.getScalarQuantizer().getBits() == 4
+          && similarityType == VectorSimilarityFunction.EUCLIDEAN
+          // Indicates that the vector is compressed as the byte length is not equal to the
+          // dimension count
+          && vectorValues.getVectorByteLength() != vectorValues.dimension()) {
+        return delegate.getRandomVectorScorer(similarityType, vectorValues, queryVector);
+      }
+      checkDimensions(queryVector.length, vectorValues.dimension());
+      var scalarQuantizer = quantizedByteVectorValues.getScalarQuantizer();
+      byte[] targetBytes = new byte[queryVector.length];
+      float offsetCorrection =
+          quantizeQuery(queryVector, targetBytes, similarityType, scalarQuantizer);
+      var scalarScorer =
+          Lucene99MemorySegmentScalarQuantizedVectorScorer.create(
+              similarityType,
+              targetBytes,
+              offsetCorrection,
+              scalarQuantizer.getConstantMultiplier(),
+              scalarQuantizer.getBits(),
+              quantizedByteVectorValues);
+      if (scalarScorer.isPresent()) {
+        return scalarScorer.get();
+      }
+    }
+    return delegate.getRandomVectorScorer(similarityType, vectorValues, queryVector);
+  }
+
+  @Override
+  public RandomVectorScorer getRandomVectorScorer(
+      VectorSimilarityFunction similarityType,
+      RandomAccessVectorValues vectorValues,
+      byte[] queryVector)
+      throws IOException {
+    return delegate.getRandomVectorScorer(similarityType, vectorValues, queryVector);
+  }
+
+  static void checkDimensions(int queryLen, int fieldLen) {
+    if (queryLen != fieldLen) {
+      throw new IllegalArgumentException(
+          "vector query dimension: " + queryLen + " differs from field dimension: " + fieldLen);
+    }
+  }
+
+  @Override
+  public String toString() {
+    return "Lucene99MemorySegmentScalarQuantizedFlatVectorsScorer()";
+  }
+}