Implement off-heap quantized scoring

Author: Kaival Parikh

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

@@ -37,4 +37,8 @@ private FlatVectorScorerUtil() {}
   public static FlatVectorsScorer getLucene99FlatVectorsScorer() {
     return IMPL.getLucene99FlatVectorsScorer();
   }
+
+  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;
@@ -70,7 +70,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() {
@@ -117,8 +117,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(DefaultFlatVectorScorer.INSTANCE);
+  }
+
   @Override
   public PostingDecodingUtil newPostingDecodingUtil(IndexInput input) {
     return new PostingDecodingUtil(input);

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

@@ -111,6 +111,9 @@ public static VectorizationProvider getInstance() {
   /** Returns a FlatVectorsScorer that supports the Lucene99 format. */
   public abstract FlatVectorsScorer getLucene99FlatVectorsScorer();
 
+  /** Returns a FlatVectorsScorer that supports 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/java24/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentScalarQuantizedVectorScorer.java

@@ -0,0 +1,319 @@
+/*
+ * 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 java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static org.apache.lucene.codecs.hnsw.ScalarQuantizedVectorScorer.quantizeQuery;
+
+import java.io.IOException;
+import java.lang.foreign.MemorySegment;
+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.index.KnnVectorValues;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.store.MemorySegmentAccessInput;
+import org.apache.lucene.util.VectorUtil;
+import org.apache.lucene.util.hnsw.RandomVectorScorer;
+import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier;
+import org.apache.lucene.util.hnsw.UpdateableRandomVectorScorer;
+import org.apache.lucene.util.quantization.QuantizedByteVectorValues;
+import org.apache.lucene.util.quantization.ScalarQuantizer;
+
+class Lucene99MemorySegmentScalarQuantizedVectorScorer implements FlatVectorsScorer {
+  static final Lucene99MemorySegmentScalarQuantizedVectorScorer INSTANCE =
+      new Lucene99MemorySegmentScalarQuantizedVectorScorer();
+
+  private static final FlatVectorsScorer DELEGATE =
+      new Lucene99ScalarQuantizedVectorScorer(DefaultFlatVectorScorer.INSTANCE);
+
+  private Lucene99MemorySegmentScalarQuantizedVectorScorer() {}
+
+  @Override
+  public RandomVectorScorerSupplier getRandomVectorScorerSupplier(
+      VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues)
+      throws IOException {
+    if (vectorValues instanceof QuantizedByteVectorValues quantized
+        && quantized.getSlice() instanceof MemorySegmentAccessInput input) {
+      return new RandomVectorScorerSupplierImpl(similarityFunction, quantized, input);
+    }
+    return DELEGATE.getRandomVectorScorerSupplier(similarityFunction, vectorValues);
+  }
+
+  @Override
+  public RandomVectorScorer getRandomVectorScorer(
+      VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues, float[] target)
+      throws IOException {
+    if (vectorValues instanceof QuantizedByteVectorValues quantized
+        && quantized.getSlice() instanceof MemorySegmentAccessInput input) {
+      return new RandomVectorScorerImpl(similarityFunction, quantized, input, target);
+    }
+    return DELEGATE.getRandomVectorScorer(similarityFunction, vectorValues, target);
+  }
+
+  @Override
+  public RandomVectorScorer getRandomVectorScorer(
+      VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues, byte[] target)
+      throws IOException {
+    return DELEGATE.getRandomVectorScorer(similarityFunction, vectorValues, target);
+  }
+
+  private abstract static class RandomVectorScorerBase
+      extends RandomVectorScorer.AbstractRandomVectorScorer {
+
+    private final ScalarQuantizer quantizer;
+    private final float constMultiplier;
+    private final MemorySegmentAccessInput input;
+    private final int vectorByteSize;
+    private final int nodeSize;
+    private final Scorer scorer;
+    private final Scaler scaler;
+    private byte[] scratch;
+
+    RandomVectorScorerBase(
+        VectorSimilarityFunction similarityFunction,
+        QuantizedByteVectorValues values,
+        MemorySegmentAccessInput input) {
+      super(values);
+
+      this.quantizer = values.getScalarQuantizer();
+      this.constMultiplier = this.quantizer.getConstantMultiplier();
+      this.input = input;
+      this.vectorByteSize = values.getVectorByteLength();
+      this.nodeSize = this.vectorByteSize + Float.BYTES;
+
+      this.scorer =
+          switch (similarityFunction) {
+            case EUCLIDEAN -> {
+              if (this.quantizer.getBits() <= 4) {
+                if (this.vectorByteSize != values.dimension()) {
+                  yield this::compressedInt4Euclidean;
+                }
+                yield this::int4Euclidean;
+              }
+              yield this::euclidean;
+            }
+            case DOT_PRODUCT, COSINE, MAXIMUM_INNER_PRODUCT -> {
+              if (this.quantizer.getBits() <= 4) {
+                if (this.vectorByteSize != values.dimension()) {
+                  yield this::compressedInt4DotProduct;
+                }
+                yield this::int4DotProduct;
+              }
+              yield this::dotProduct;
+            }
+          };
+
+      this.scaler =
+          switch (similarityFunction) {
+            case EUCLIDEAN -> distance -> 1 / (1 + distance);
+            case DOT_PRODUCT, COSINE -> score -> Math.max((1 + score) / 2, 0);
+            case MAXIMUM_INNER_PRODUCT -> VectorUtil::scaleMaxInnerProductScore;
+          };
+
+      checkInvariants();
+    }
+
+    final void checkInvariants() {
+      if (input.length() < (long) nodeSize * maxOrd()) {
+        throw new IllegalArgumentException("input length is less than expected vector data");
+      }
+    }
+
+    final void checkOrdinal(int ord) {
+      if (ord < 0 || ord >= maxOrd()) {
+        throw new IllegalArgumentException("illegal ordinal: " + ord);
+      }
+    }
+
+    ScalarQuantizer getQuantizer() {
+      return quantizer;
+    }
+
+    @SuppressWarnings("restricted")
+    Node getNode(int ord) throws IOException {
+      checkOrdinal(ord);
+      long byteOffset = (long) ord * nodeSize;
+      MemorySegment node = input.segmentSliceOrNull(byteOffset, nodeSize);
+      if (node == null) {
+        if (scratch == null) {
+          scratch = new byte[nodeSize];
+        }
+        input.readBytes(byteOffset, scratch, 0, nodeSize);
+        node = MemorySegment.ofArray(scratch);
+      }
+      return new Node(node.reinterpret(vectorByteSize), node.get(JAVA_FLOAT, vectorByteSize));
+    }
+
+    float scoreBody(int ord, float queryOffset) throws IOException {
+      checkOrdinal(ord);
+      Node node = getNode(ord);
+      return scaler.scale(scorer.score(node.vector) * constMultiplier + node.offset + queryOffset);
+    }
+
+    abstract int euclidean(MemorySegment doc);
+
+    abstract int int4Euclidean(MemorySegment doc);
+
+    abstract int compressedInt4Euclidean(MemorySegment doc);
+
+    abstract int dotProduct(MemorySegment doc);
+
+    abstract int int4DotProduct(MemorySegment doc);
+
+    abstract int compressedInt4DotProduct(MemorySegment doc);
+
+    record Node(MemorySegment vector, float offset) {}
+
+    @FunctionalInterface
+    private interface Scorer {
+      int score(MemorySegment doc) throws IOException;
+    }
+
+    @FunctionalInterface
+    private interface Scaler {
+      float scale(float score);
+    }
+  }
+
+  private static class RandomVectorScorerImpl extends RandomVectorScorerBase {
+    private final byte[] targetBytes;
+    private final float queryOffset;
+
+    RandomVectorScorerImpl(
+        VectorSimilarityFunction similarityFunction,
+        QuantizedByteVectorValues values,
+        MemorySegmentAccessInput input,
+        float[] target) {
+      super(similarityFunction, values, input);
+      this.targetBytes = new byte[target.length];
+      this.queryOffset = quantizeQuery(target, targetBytes, similarityFunction, getQuantizer());
+    }
+
+    @Override
+    public float score(int node) throws IOException {
+      return scoreBody(node, queryOffset);
+    }
+
+    @Override
+    int euclidean(MemorySegment doc) {
+      return PanamaVectorUtilSupport.squareDistance(targetBytes, doc);
+    }
+
+    @Override
+    int int4Euclidean(MemorySegment doc) {
+      // TODO
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    int compressedInt4Euclidean(MemorySegment doc) {
+      // TODO
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    int dotProduct(MemorySegment doc) {
+      return PanamaVectorUtilSupport.dotProduct(targetBytes, doc);
+    }
+
+    @Override
+    int int4DotProduct(MemorySegment doc) {
+      return PanamaVectorUtilSupport.int4DotProduct(targetBytes, false, doc, false);
+    }
+
+    @Override
+    int compressedInt4DotProduct(MemorySegment doc) {
+      return PanamaVectorUtilSupport.int4DotProduct(targetBytes, false, doc, true);
+    }
+  }
+
+  private record RandomVectorScorerSupplierImpl(
+      VectorSimilarityFunction similarityFunction,
+      QuantizedByteVectorValues values,
+      MemorySegmentAccessInput input)
+      implements RandomVectorScorerSupplier {
+
+    @Override
+    public UpdateableRandomVectorScorer scorer() {
+      return new UpdateableRandomVectorScorerImpl(similarityFunction, values, input);
+    }
+
+    @Override
+    public RandomVectorScorerSupplier copy() {
+      return new RandomVectorScorerSupplierImpl(similarityFunction, values, input);
+    }
+  }
+
+  private static class UpdateableRandomVectorScorerImpl extends RandomVectorScorerBase
+      implements UpdateableRandomVectorScorer {
+    private MemorySegment query;
+    private float queryOffset;
+
+    UpdateableRandomVectorScorerImpl(
+        VectorSimilarityFunction similarityFunction,
+        QuantizedByteVectorValues values,
+        MemorySegmentAccessInput input) {
+      super(similarityFunction, values, input);
+    }
+
+    @Override
+    public void setScoringOrdinal(int ord) throws IOException {
+      checkOrdinal(ord);
+      Node node = getNode(ord);
+      query = node.vector;
+      queryOffset = node.offset;
+    }
+
+    @Override
+    public float score(int node) throws IOException {
+      return scoreBody(node, queryOffset);
+    }
+
+    @Override
+    int euclidean(MemorySegment doc) {
+      return PanamaVectorUtilSupport.squareDistance(query, doc);
+    }
+
+    @Override
+    int int4Euclidean(MemorySegment doc) {
+      // TODO
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    int compressedInt4Euclidean(MemorySegment doc) {
+      // TODO
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    int dotProduct(MemorySegment doc) {
+      return PanamaVectorUtilSupport.dotProduct(query, doc);
+    }
+
+    @Override
+    int int4DotProduct(MemorySegment doc) {
+      return PanamaVectorUtilSupport.int4DotProduct(query, false, doc, false);
+    }
+
+    @Override
+    int compressedInt4DotProduct(MemorySegment doc) {
+      return PanamaVectorUtilSupport.int4DotProduct(query, true, doc, true);
+    }
+  }
+}