Add support for uint8 distance comparison (#15148)

Add distance functions that treat `byte[]` as unsigned and use them in `ScalarQuantizer` code paths. `ScalarQuantizer`
assumes that all math will be unsigned but this can't be true when all 8 bits of a `byte` are used. In all cases where this
is used values are widened to `short` or `int` before any operations occur so we can perform unsigned widening instead.

Fixes are required in a few other places to support unsigned extension, in particular this is required when recalculating
the vector offset as part of merging.

This will be useful for supporting 8 bit OSQ in #15064

Author: mccullocht

lucene/CHANGES.txt

@@ -123,7 +123,7 @@ New Features
 
 Improvements
 ---------------------
-(No changes)
+# GITHUB#15148: Add support uint8 distance and allow 8 bit scalar quantization (Trevor McCulloch)
 
 Optimizations
 ---------------------

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

@@ -119,6 +119,17 @@ public int binaryDotProductVector() {
     return VectorUtil.dotProduct(bytesA, bytesB);
   }
 
+  @Benchmark
+  public int binaryDotProductUint8Scalar() {
+    return VectorUtil.uint8DotProduct(bytesA, bytesB);
+  }
+
+  @Benchmark
+  @Fork(jvmArgsPrepend = {"--add-modules=jdk.incubator.vector"})
+  public int binaryDotProductUint8Vector() {
+    return VectorUtil.uint8DotProduct(bytesA, bytesB);
+  }
+
   @Benchmark
   public int binarySquareScalar() {
     return VectorUtil.squareDistance(bytesA, bytesB);
@@ -130,6 +141,17 @@ public int binarySquareVector() {
     return VectorUtil.squareDistance(bytesA, bytesB);
   }
 
+  @Benchmark
+  public int binarySquareUint8Scalar() {
+    return VectorUtil.uint8SquareDistance(bytesA, bytesB);
+  }
+
+  @Benchmark
+  @Fork(jvmArgsPrepend = {"--add-modules=jdk.incubator.vector"})
+  public int binarySquareUint8Vector() {
+    return VectorUtil.uint8SquareDistance(bytesA, bytesB);
+  }
+
   @Benchmark
   public int binaryHalfByteScalar() {
     return VectorUtil.int4DotProduct(halfBytesA, halfBytesB);

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

@@ -156,7 +156,7 @@ private Euclidean(QuantizedByteVectorValues values, float constMultiplier, byte[
     @Override
     public float score(int node) throws IOException {
       byte[] nodeVector = values.vectorValue(node);
-      int squareDistance = VectorUtil.squareDistance(nodeVector, targetBytes);
+      int squareDistance = VectorUtil.uint8SquareDistance(nodeVector, targetBytes);
       float adjustedDistance = squareDistance * constMultiplier;
       return 1 / (1f + adjustedDistance);
     }
@@ -194,8 +194,9 @@ public DotProduct(
     public float score(int vectorOrdinal) throws IOException {
       byte[] storedVector = values.vectorValue(vectorOrdinal);
       float vectorOffset = values.getScoreCorrectionConstant(vectorOrdinal);
-      int dotProduct = VectorUtil.dotProduct(storedVector, targetBytes);
-      // For the current implementation of scalar quantization, all dotproducts should be >= 0;
+      int dotProduct = VectorUtil.uint8DotProduct(storedVector, targetBytes);
+      // For the current implementation of scalar quantization, all dotproducts should
+      // be >= 0;
       assert dotProduct >= 0;
       float adjustedDistance = dotProduct * constMultiplier + offsetCorrection + vectorOffset;
       return scoreAdjustmentFunction.apply(adjustedDistance);
@@ -208,9 +209,10 @@ public void setScoringOrdinal(int node) throws IOException {
     }
   }
 
-  // TODO consider splitting this into two classes. right now the "query" vector is always
+  // TODO consider splitting this into two classes. right now the "query" vector
+  // is always
   // decompressed
-  //    it could stay compressed if we had a compressed version of the target vector
+  // it could stay compressed if we had a compressed version of the target vector
   private static class CompressedInt4DotProduct
       extends UpdateableRandomVectorScorer.AbstractUpdateableRandomVectorScorer {
     private final float constMultiplier;
@@ -237,13 +239,15 @@ private CompressedInt4DotProduct(
 
     @Override
     public float score(int vectorOrdinal) throws IOException {
-      // get compressed vector, in Lucene99, vector values are stored and have a single value for
+      // get compressed vector, in Lucene99, vector values are stored and have a
+      // single value for
       // offset correction
       values.getSlice().seek((long) vectorOrdinal * (values.getVectorByteLength() + Float.BYTES));
       values.getSlice().readBytes(compressedVector, 0, compressedVector.length);
       float vectorOffset = values.getScoreCorrectionConstant(vectorOrdinal);
       int dotProduct = VectorUtil.int4DotProductPacked(targetBytes, compressedVector);
-      // For the current implementation of scalar quantization, all dotproducts should be >= 0;
+      // For the current implementation of scalar quantization, all dotproducts should
+      // be >= 0;
       assert dotProduct >= 0;
       float adjustedDistance = dotProduct * constMultiplier + offsetCorrection + vectorOffset;
       return scoreAdjustmentFunction.apply(adjustedDistance);
@@ -283,7 +287,8 @@ public float score(int vectorOrdinal) throws IOException {
       byte[] storedVector = values.vectorValue(vectorOrdinal);
       float vectorOffset = values.getScoreCorrectionConstant(vectorOrdinal);
       int dotProduct = VectorUtil.int4DotProduct(storedVector, targetBytes);
-      // For the current implementation of scalar quantization, all dotproducts should be >= 0;
+      // For the current implementation of scalar quantization, all dotproducts should
+      // be >= 0;
       assert dotProduct >= 0;
       float adjustedDistance = dotProduct * constMultiplier + offsetCorrection + vectorOffset;
       return scoreAdjustmentFunction.apply(adjustedDistance);

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

@@ -34,10 +34,8 @@
 public class Lucene99ScalarQuantizedVectorsFormat extends FlatVectorsFormat {
 
   // The bits that are allowed for scalar quantization
-  // We only allow signed byte (7), and half-byte (4)
-  // NOTE: we used to allow 8 bits as well, but it was broken so we removed it
-  // (https://github.com/apache/lucene/issues/13519)
-  private static final int ALLOWED_BITS = (1 << 7) | (1 << 4);
+  // We only allow unsigned byte (8), signed byte (7), and half-byte (4)
+  private static final int ALLOWED_BITS = (1 << 8) | (1 << 7) | (1 << 4);
   public static final String QUANTIZED_VECTOR_COMPONENT = "QVEC";
 
   public static final String NAME = "Lucene99ScalarQuantizedVectorsFormat";

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

@@ -154,6 +154,15 @@ public int dotProduct(byte[] a, byte[] b) {
     return total;
   }
 
+  @Override
+  public int uint8DotProduct(byte[] a, byte[] b) {
+    int total = 0;
+    for (int i = 0; i < a.length; i++) {
+      total += Byte.toUnsignedInt(a[i]) * Byte.toUnsignedInt(b[i]);
+    }
+    return total;
+  }
+
   @Override
   public int int4DotProduct(byte[] a, boolean apacked, byte[] b, boolean bpacked) {
     assert (apacked && bpacked) == false;
@@ -201,6 +210,17 @@ public int squareDistance(byte[] a, byte[] b) {
     return squareSum;
   }
 
+  @Override
+  public int uint8SquareDistance(byte[] a, byte[] b) {
+    // Note: this will not overflow if dim < 2^16, since max(ubyte * ubyte) = 2^16.
+    int squareSum = 0;
+    for (int i = 0; i < a.length; i++) {
+      int diff = Byte.toUnsignedInt(a[i]) - Byte.toUnsignedInt(b[i]);
+      squareSum += diff * diff;
+    }
+    return squareSum;
+  }
+
   @Override
   public int findNextGEQ(int[] buffer, int target, int from, int to) {
     for (int i = from; i < to; ++i) {
@@ -281,7 +301,7 @@ float recalculateOffset(byte[] vector, int start, float oldAlpha, float oldMinQu
       float correction = 0;
       for (int i = start; i < vector.length; i++) {
         // undo the old quantization
-        float v = (oldAlpha * vector[i]) + oldMinQuantile;
+        float v = (oldAlpha * Byte.toUnsignedInt(vector[i])) + oldMinQuantile;
         correction += quantizeFloat(v, null, 0);
       }
       return correction;

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

@@ -36,6 +36,9 @@ public interface VectorUtilSupport {
   /** Returns the dot product computed over signed bytes. */
   int dotProduct(byte[] a, byte[] b);
 
+  /** Returns the dot product computed as though the bytes were unsigned. */
+  int uint8DotProduct(byte[] a, byte[] b);
+
   /** Returns the dot product over the computed bytes, assuming the values are int4 encoded. */
   int int4DotProduct(byte[] a, boolean apacked, byte[] b, boolean bpacked);
 
@@ -45,6 +48,9 @@ public interface VectorUtilSupport {
   /** Returns the sum of squared differences of the two byte vectors. */
   int squareDistance(byte[] a, byte[] b);
 
+  /** Returns the sum of squared differences of the two unsigned byte vectors. */
+  int uint8SquareDistance(byte[] a, byte[] b);
+
   /**
    * Given an array {@code buffer} that is sorted between indexes {@code 0} inclusive and {@code to}
    * exclusive, find the first array index whose value is greater than or equal to {@code target}.

lucene/core/src/java/org/apache/lucene/util/VectorUtil.java

@@ -113,6 +113,14 @@ public static int squareDistance(byte[] a, byte[] b) {
     return IMPL.squareDistance(a, b);
   }
 
+  /** Returns the sum of squared differences of the two vectors where each byte is unsigned */
+  public static int uint8SquareDistance(byte[] a, byte[] b) {
+    if (a.length != b.length) {
+      throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);
+    }
+    return IMPL.uint8SquareDistance(a, b);
+  }
+
   /**
    * Modifies the argument to be unit length, dividing by its l2-norm. IllegalArgumentException is
    * thrown for zero vectors.
@@ -167,6 +175,20 @@ public static int dotProduct(byte[] a, byte[] b) {
     return IMPL.dotProduct(a, b);
   }
 
+  /**
+   * Dot product over bytes assuming that the values are actually unsigned.
+   *
+   * @param a uint8 byte vector
+   * @param b another uint8 byte vector of the same dimension
+   * @return the value of the dot product of the two vectors
+   */
+  public static int uint8DotProduct(byte[] a, byte[] b) {
+    if (a.length != b.length) {
+      throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);
+    }
+    return IMPL.uint8DotProduct(a, b);
+  }
+
   public static int int4DotProduct(byte[] a, byte[] b) {
     if (a.length != b.length) {
       throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);

lucene/core/src/java/org/apache/lucene/util/quantization/ScalarQuantizedVectorSimilarity.java

@@ -42,10 +42,12 @@ static ScalarQuantizedVectorSimilarity fromVectorSimilarity(
       case EUCLIDEAN -> new Euclidean(constMultiplier);
       case COSINE, DOT_PRODUCT ->
           new DotProduct(
-              constMultiplier, bits <= 4 ? VectorUtil::int4DotProduct : VectorUtil::dotProduct);
+              constMultiplier,
+              bits <= 4 ? VectorUtil::int4DotProduct : VectorUtil::uint8DotProduct);
       case MAXIMUM_INNER_PRODUCT ->
           new MaximumInnerProduct(
-              constMultiplier, bits <= 4 ? VectorUtil::int4DotProduct : VectorUtil::dotProduct);
+              constMultiplier,
+              bits <= 4 ? VectorUtil::int4DotProduct : VectorUtil::uint8DotProduct);
     };
   }
 
@@ -62,7 +64,7 @@ public Euclidean(float constMultiplier) {
     @Override
     public float score(
         byte[] queryVector, float queryVectorOffset, byte[] storedVector, float vectorOffset) {
-      int squareDistance = VectorUtil.squareDistance(storedVector, queryVector);
+      int squareDistance = VectorUtil.uint8SquareDistance(storedVector, queryVector);
       float adjustedDistance = squareDistance * constMultiplier;
       return 1 / (1f + adjustedDistance);
     }

lucene/core/src/java/org/apache/lucene/util/quantization/ScalarQuantizer.java

@@ -165,7 +165,7 @@ public float recalculateCorrectiveOffset(
   public void deQuantize(byte[] src, float[] dest) {
     assert src.length == dest.length;
     for (int i = 0; i < src.length; i++) {
-      dest[i] = (alpha * src[i]) + minQuantile;
+      dest[i] = (alpha * Byte.toUnsignedInt(src[i])) + minQuantile;
     }
   }