Implement off-heap scoring for OSQ 4, 7, and 8 bit representations

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

@@ -17,6 +17,7 @@
 
 package org.apache.lucene.codecs.hnsw;
 
+import org.apache.lucene.codecs.lucene104.AsymmetricScalarQuantizeFlatVectorsScorer;
 import org.apache.lucene.internal.vectorization.VectorizationProvider;
 
 /**
@@ -41,4 +42,14 @@ public static FlatVectorsScorer getLucene99FlatVectorsScorer() {
   public static FlatVectorsScorer getLucene99ScalarQuantizedVectorsScorer() {
     return IMPL.getLucene99ScalarQuantizedVectorsScorer();
   }
+
+  /**
+   * Returns a FlatVectorsScorer that supports the Lucene104 scalar quantized format. Scorers
+   * retrieved through this method may be optimized on certain platforms. Otherwise, a
+   * DefaultFlatVectorScorer is returned.
+   */
+  public static AsymmetricScalarQuantizeFlatVectorsScorer
+      getLucene104ScalarQuantizedFlatVectorsScorer() {
+    return IMPL.getLucene104ScalarQuantizedVectorsScorer();
+  }
 }

lucene/core/src/java/org/apache/lucene/codecs/lucene104/AsymmetricScalarQuantizeFlatVectorsScorer.java

@@ -0,0 +1,51 @@
+/*
+ * 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.
+ */
+
+// XXX DO NOT MERGE FIX NAME
+
+package org.apache.lucene.codecs.lucene104;
+
+import java.io.IOException;
+import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier;
+
+/**
+ * Extension of {@link FlatVectorsScorer} that allows using two different vector codings for the
+ * "scoring" or "query" vectors and the "target" or "doc" vectors.
+ */
+public interface AsymmetricScalarQuantizeFlatVectorsScorer extends FlatVectorsScorer {
+  /**
+   * Returns a {@link RandomVectorScorerSupplier} that can be used to score asymmetric vector
+   * representations, typically a higher fidelity "scoring" vector against a lower fidelity "target"
+   * vector. This is used during indexing to improve the quality of the index data structure during
+   * build/merge; only the targetVectors are saved.
+   *
+   * <p>This may only be used when ScalarEncoding.isAsymmetric().
+   *
+   * @param similarityFunction the similarity function to use
+   * @param scoringVectors higher fidelity scoring vectors to use as queries.
+   * @param targetVectors lower fidelity vectors to use as documents.
+   * @return a {@link RandomVectorScorerSupplier} that can be used to score vectors
+   * @throws IOException if an I/O error occurs
+   */
+  RandomVectorScorerSupplier getRandomVectorScorerSupplier(
+      VectorSimilarityFunction similarityFunction,
+      QuantizedByteVectorValues scoringVectors,
+      QuantizedByteVectorValues targetVectors)
+      throws IOException;
+}

lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorScorer.java

@@ -22,6 +22,7 @@
 
 import java.io.IOException;
 import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
+import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding;
 import org.apache.lucene.index.KnnVectorValues;
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.util.ArrayUtil;
@@ -32,7 +33,8 @@
 import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
 
 /** Vector scorer over OptimizedScalarQuantized vectors */
-public class Lucene104ScalarQuantizedVectorScorer implements FlatVectorsScorer {
+public class Lucene104ScalarQuantizedVectorScorer
+    implements AsymmetricScalarQuantizeFlatVectorsScorer {
   private final FlatVectorsScorer nonQuantizedDelegate;
 
   public Lucene104ScalarQuantizedVectorScorer(FlatVectorsScorer nonQuantizedDelegate) {
@@ -106,7 +108,8 @@ public RandomVectorScorer getRandomVectorScorer(
     return nonQuantizedDelegate.getRandomVectorScorer(similarityFunction, vectorValues, target);
   }
 
-  RandomVectorScorerSupplier getRandomVectorScorerSupplier(
+  @Override
+  public RandomVectorScorerSupplier getRandomVectorScorerSupplier(
       VectorSimilarityFunction similarityFunction,
       QuantizedByteVectorValues scoringVectors,
       QuantizedByteVectorValues targetVectors) {
@@ -239,7 +242,7 @@ private static float quantizedScore(
       int targetOrd,
       VectorSimilarityFunction similarityFunction)
       throws IOException {
-    var scalarEncoding = targetVectors.getScalarEncoding();
+    ScalarEncoding scalarEncoding = targetVectors.getScalarEncoding();
     byte[] quantizedDoc = targetVectors.vectorValue(targetOrd);
     float qcDist =
         switch (scalarEncoding) {
@@ -249,8 +252,32 @@ private static float quantizedScore(
           case SINGLE_BIT_QUERY_NIBBLE ->
               VectorUtil.int4BitDotProduct(quantizedQuery, quantizedDoc);
         };
-    OptimizedScalarQuantizer.QuantizationResult indexCorrections =
-        targetVectors.getCorrectiveTerms(targetOrd);
+    return quantizedScore(
+        similarityFunction,
+        targetVectors,
+        qcDist,
+        queryCorrections,
+        targetVectors.getCorrectiveTerms(targetOrd));
+  }
+
+  /**
+   * Transforms the dotProduct of a query and index vector into a score.
+   *
+   * @param similarityFunction similarity function used to compute the score
+   * @param targetVectors target vector set; used for metadata
+   * @param dotProduct dot product of query and index vectors.
+   * @param queryCorrections corrective terms for the query vector
+   * @param indexCorrections corrective terms for the index vector
+   * @return a score value greater than or equal to 0
+   */
+  public static float quantizedScore(
+      VectorSimilarityFunction similarityFunction,
+      QuantizedByteVectorValues targetVectors,
+      float dotProduct,
+      OptimizedScalarQuantizer.QuantizationResult queryCorrections,
+      OptimizedScalarQuantizer.QuantizationResult indexCorrections)
+      throws IOException {
+    ScalarEncoding scalarEncoding = targetVectors.getScalarEncoding();
     float queryScale = SCALE_LUT[scalarEncoding.getQueryBits() - 1];
     float scale = SCALE_LUT[scalarEncoding.getBits() - 1];
     float x1 = indexCorrections.quantizedComponentSum();
@@ -261,7 +288,7 @@ private static float quantizedScore(
     float ly = (queryCorrections.upperInterval() - ay) * queryScale;
     float y1 = queryCorrections.quantizedComponentSum();
     float score =
-        ax * ay * targetVectors.dimension() + ay * lx * x1 + ax * ly * y1 + lx * ly * qcDist;
+        ax * ay * targetVectors.dimension() + ay * lx * x1 + ax * ly * y1 + lx * ly * dotProduct;
     // For euclidean, we need to invert the score and apply the additional correction, which is
     // assumed to be the squared l2norm of the centroid centered vectors.
     if (similarityFunction == EUCLIDEAN) {

lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsFormat.java

@@ -105,8 +105,8 @@ public class Lucene104ScalarQuantizedVectorsFormat extends FlatVectorsFormat {
   private static final FlatVectorsFormat rawVectorFormat =
       new Lucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer());
 
-  private static final Lucene104ScalarQuantizedVectorScorer scorer =
-      new Lucene104ScalarQuantizedVectorScorer(FlatVectorScorerUtil.getLucene99FlatVectorsScorer());
+  private static final AsymmetricScalarQuantizeFlatVectorsScorer scorer =
+      FlatVectorScorerUtil.getLucene104ScalarQuantizedFlatVectorsScorer();
 
   private final ScalarEncoding encoding;
 

lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsReader.java

@@ -27,6 +27,7 @@
 import org.apache.lucene.codecs.CodecUtil;
 import org.apache.lucene.codecs.KnnVectorsReader;
 import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
+import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
 import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding;
 import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration;
 import org.apache.lucene.index.ByteVectorValues;
@@ -66,12 +67,10 @@ class Lucene104ScalarQuantizedVectorsReader extends FlatVectorsReader
   private final Map<String, FieldEntry> fields = new HashMap<>();
   private final IndexInput quantizedVectorData;
   private final FlatVectorsReader rawVectorsReader;
-  private final Lucene104ScalarQuantizedVectorScorer vectorScorer;
+  private final FlatVectorsScorer vectorScorer;
 
   Lucene104ScalarQuantizedVectorsReader(
-      SegmentReadState state,
-      FlatVectorsReader rawVectorsReader,
-      Lucene104ScalarQuantizedVectorScorer vectorsScorer)
+      SegmentReadState state, FlatVectorsReader rawVectorsReader, FlatVectorsScorer vectorsScorer)
       throws IOException {
     super(vectorsScorer);
     this.vectorScorer = vectorsScorer;

lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java

@@ -69,7 +69,7 @@ public class Lucene104ScalarQuantizedVectorsWriter extends FlatVectorsWriter {
   private final IndexOutput meta, vectorData;
   private final ScalarEncoding encoding;
   private final FlatVectorsWriter rawVectorDelegate;
-  private final Lucene104ScalarQuantizedVectorScorer vectorsScorer;
+  private final AsymmetricScalarQuantizeFlatVectorsScorer vectorsScorer;
   private boolean finished;
 
   /**
@@ -81,7 +81,7 @@ protected Lucene104ScalarQuantizedVectorsWriter(
       SegmentWriteState state,
       ScalarEncoding encoding,
       FlatVectorsWriter rawVectorDelegate,
-      Lucene104ScalarQuantizedVectorScorer vectorsScorer)
+      AsymmetricScalarQuantizeFlatVectorsScorer vectorsScorer)
       throws IOException {
     super(vectorsScorer);
     this.encoding = encoding;
@@ -831,6 +831,11 @@ public int dimension() {
       return values.dimension();
     }
 
+    @Override
+    public IndexInput getSlice() {
+      return null;
+    }
+
     @Override
     public OptimizedScalarQuantizer getQuantizer() {
       throw new UnsupportedOperationException();

lucene/core/src/java/org/apache/lucene/codecs/lucene104/OffHeapScalarQuantizedVectorValues.java

@@ -169,7 +169,12 @@ public float[] getCentroid() {
 
   @Override
   public int getVectorByteLength() {
-    return dimension;
+    return this.encoding.getDocPackedLength(dimension);
+  }
+
+  @Override
+  public IndexInput getSlice() {
+    return slice;
   }
 
   static void packNibbles(byte[] unpacked, byte[] packed) {

lucene/core/src/java/org/apache/lucene/codecs/lucene104/QuantizedByteVectorValues.java

@@ -18,12 +18,13 @@
 
 import java.io.IOException;
 import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding;
+import org.apache.lucene.codecs.lucene95.HasIndexSlice;
 import org.apache.lucene.index.ByteVectorValues;
 import org.apache.lucene.search.VectorScorer;
 import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
 
 /** Scalar quantized byte vector values */
-abstract class QuantizedByteVectorValues extends ByteVectorValues {
+public abstract class QuantizedByteVectorValues extends ByteVectorValues implements HasIndexSlice {
 
   /**
    * Retrieve the corrective terms for the given vector ordinal. For the dot-product family of

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

@@ -19,6 +19,8 @@
 
 import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer;
 import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
+import org.apache.lucene.codecs.lucene104.AsymmetricScalarQuantizeFlatVectorsScorer;
+import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorScorer;
 import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorScorer;
 import org.apache.lucene.store.IndexInput;
 
@@ -46,6 +48,11 @@ public FlatVectorsScorer getLucene99ScalarQuantizedVectorsScorer() {
     return new Lucene99ScalarQuantizedVectorScorer(DefaultFlatVectorScorer.INSTANCE);
   }
 
+  @Override
+  public AsymmetricScalarQuantizeFlatVectorsScorer getLucene104ScalarQuantizedVectorsScorer() {
+    return new Lucene104ScalarQuantizedVectorScorer(DefaultFlatVectorScorer.INSTANCE);
+  }
+
   @Override
   public PostingDecodingUtil newPostingDecodingUtil(IndexInput input) {
     return new PostingDecodingUtil(input);

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

@@ -29,6 +29,7 @@
 import java.util.logging.Logger;
 import java.util.stream.Stream;
 import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
+import org.apache.lucene.codecs.lucene104.AsymmetricScalarQuantizeFlatVectorsScorer;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.util.Constants;
 import org.apache.lucene.util.VectorUtil;
@@ -112,6 +113,10 @@ public static VectorizationProvider getInstance() {
   /** Returns a FlatVectorsScorer that supports the Lucene99 format. */
   public abstract FlatVectorsScorer getLucene99ScalarQuantizedVectorsScorer();
 
+  /** Returns a FlatVectorsScorer that supports the Lucene104 quantized format. */
+  public abstract AsymmetricScalarQuantizeFlatVectorsScorer
+      getLucene104ScalarQuantizedVectorsScorer();
+
   /** Create a new {@link PostingDecodingUtil} for the given {@link IndexInput}. */
   public abstract PostingDecodingUtil newPostingDecodingUtil(IndexInput input) throws IOException;
 
@@ -136,7 +141,8 @@ static VectorizationProvider lookup(boolean testMode) {
             "Java runtime is using JVMCI Compiler; Java vector incubator API can't be enabled.");
         return new DefaultVectorizationProvider();
       }
-      // is the incubator module present and readable (JVM providers may to exclude them or it is
+      // is the incubator module present and readable (JVM providers may to exclude
+      // them or it is
       // build with jlink)
       final var vectorMod = lookupVectorModule();
       if (vectorMod.isEmpty()) {
@@ -158,7 +164,8 @@ static VectorizationProvider lookup(boolean testMode) {
         }
       }
       try {
-        // we use method handles with lookup, so we do not need to deal with setAccessible as we
+        // we use method handles with lookup, so we do not need to deal with
+        // setAccessible as we
         // have private access through the lookup:
         final var lookup = MethodHandles.lookup();
         final var cls =