iter

Author: benwtrent

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

@@ -410,8 +410,8 @@ static DocsWithFieldSet writeBinarizedVectorAndQueryData(
         Math.max(
             encoding.getDiscreteDimensions(floatVectorValues.dimension()),
             queryEncoding.getDiscreteDimensions(floatVectorValues.dimension()));
-    assert discretizedDims % encoding.getBits() == 0;
-    assert discretizedDims % queryEncoding.getBits() == 0;
+    assert discretizedDims % (8 / encoding.getBitsPerDim()) == 0;
+    assert discretizedDims % (8 / queryEncoding.getBitsPerDim()) == 0;
     byte[][] quantizationScratch = new byte[2][];
     quantizationScratch[0] = new byte[discretizedDims];
     quantizationScratch[1] = new byte[discretizedDims];
@@ -563,7 +563,9 @@ private CloseableRandomVectorScorerSupplier mergeOneFieldToIndex(
       if (finalQuantizedScoreDataInput != null) {
         scoreVectorValues =
             new OffHeapScalarQuantizedVectorValues.DenseOffHeapVectorValues(
-                fieldInfo.getVectorDimension(),
+                // pass the commonly discretized dimension to ensure both vectorValues and
+                // scoreVectorValues have the same discretized dimension
+                vectorValues.discretizedDimension(),
                 docsWithField.cardinality(),
                 centroid,
                 cDotC,
@@ -579,7 +581,7 @@ private CloseableRandomVectorScorerSupplier mergeOneFieldToIndex(
               ? vectorsScorer.getRandomVectorScorerSupplier(
                   fieldInfo.getVectorSimilarityFunction(), vectorValues)
               : vectorsScorer.getRandomVectorScorerSupplier(
-                  fieldInfo.getVectorSimilarityFunction(), vectorValues, scoreVectorValues);
+                  fieldInfo.getVectorSimilarityFunction(), scoreVectorValues, vectorValues);
       return new QuantizedCloseableRandomVectorScorerSupplier(
           scorerSupplier,
           vectorValues,

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

@@ -46,6 +46,7 @@ public abstract class OffHeapScalarQuantizedVectorValues extends QuantizedByteVe
   final byte[] vectorValue;
   final ByteBuffer byteBuffer;
   final int byteSize;
+  final int discretizedDimension;
   private int lastOrd = -1;
   final float[] correctiveValues;
   int quantizedComponentSum;
@@ -73,12 +74,14 @@ public abstract class OffHeapScalarQuantizedVectorValues extends QuantizedByteVe
     this.centroid = centroid;
     this.centroidDp = centroidDp;
     this.correctiveValues = new float[3];
-    this.byteSize = encoding.getPackedLength(dimension) + (Float.BYTES * 3) + Integer.BYTES;
-    this.byteBuffer = ByteBuffer.allocate(encoding.getPackedLength(dimension));
-    this.vectorValue = byteBuffer.array();
-    this.quantizer = quantizer;
     this.encoding = encoding;
     this.queryEncoding = queryEncoding;
+    this.discretizedDimension = calculateDiscretizedDimension(dimension, encoding, queryEncoding);
+    this.byteSize =
+        encoding.getPackedLength(discretizedDimension) + (Float.BYTES * 3) + Integer.BYTES;
+    this.byteBuffer = ByteBuffer.allocate(encoding.getPackedLength(discretizedDimension));
+    this.vectorValue = byteBuffer.array();
+    this.quantizer = quantizer;
   }
 
   @Override

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

@@ -25,6 +25,19 @@
 /** Scalar quantized byte vector values */
 abstract class QuantizedByteVectorValues extends ByteVectorValues {
 
+  public static int calculateDiscretizedDimension(
+      int dims, ScalarEncoding storageEncoding, ScalarEncoding queryEncoding) {
+    if (storageEncoding == queryEncoding) {
+      return storageEncoding.getDiscreteDimensions(dims);
+    }
+    int queryDiscretized = queryEncoding.getDiscreteDimensions(dims);
+    int docDiscretized = storageEncoding.getDiscreteDimensions(dims);
+    int maxDiscretized = Math.max(queryDiscretized, docDiscretized);
+    assert maxDiscretized % (8 / queryEncoding.getBitsPerDim()) == 0;
+    assert maxDiscretized % (8 / storageEncoding.getBitsPerDim()) == 0;
+    return maxDiscretized;
+  }
+
   /**
    * Retrieve the corrective terms for the given vector ordinal. For the dot-product family of
    * distances, the corrective terms are, in order
@@ -72,16 +85,8 @@ public abstract OptimizedScalarQuantizer.QuantizationResult getCorrectiveTerms(i
    * @return The correctly discretized dimension given the configured encodings
    */
   public int discretizedDimension() {
-    int dims = dimension();
-    if (getScalarEncoding() == getQueryScalarEncoding()) {
-      return getScalarEncoding().getDiscreteDimensions(dims);
-    }
-    int queryDiscretized = getQueryScalarEncoding().getDiscreteDimensions(dims);
-    int docDiscretized = getScalarEncoding().getDiscreteDimensions(dims);
-    int maxDiscretized = Math.max(queryDiscretized, docDiscretized);
-    assert maxDiscretized % getQueryScalarEncoding().getBits() == 0;
-    assert maxDiscretized % getScalarEncoding().getBits() == 0;
-    return maxDiscretized;
+    return calculateDiscretizedDimension(
+        dimension(), getScalarEncoding(), getQueryScalarEncoding());
   }
 
   /**

lucene/core/src/test/org/apache/lucene/codecs/lucene104/TestLucene104HnswScalarQuantizedVectorsFormat.java

@@ -30,6 +30,7 @@
 import org.apache.lucene.codecs.KnnVectorsFormat;
 import org.apache.lucene.codecs.KnnVectorsReader;
 import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding;
+import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat;
 import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader;
 import org.apache.lucene.document.Document;
 import org.apache.lucene.document.KnnFloatVectorField;
@@ -47,14 +48,36 @@
 import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase;
 import org.apache.lucene.tests.util.TestUtil;
 import org.apache.lucene.util.SameThreadExecutorService;
+import org.junit.Before;
 
 public class TestLucene104HnswScalarQuantizedVectorsFormat extends BaseKnnVectorsFormatTestCase {
 
-  private static final KnnVectorsFormat FORMAT = new Lucene104HnswScalarQuantizedVectorsFormat();
+  private KnnVectorsFormat format;
+
+  @Before
+  @Override
+  public void setUp() throws Exception {
+    var encodingValues = ScalarEncoding.values();
+    var encoding = encodingValues[random().nextInt(encodingValues.length)];
+    var queryEncoding = encoding;
+    // always assume asymmetric for now. Eventually make this more general
+    if (encoding == ScalarEncoding.SINGLE_BIT) {
+      queryEncoding = ScalarEncoding.PACKED_NIBBLE;
+    }
+    format =
+        new Lucene104HnswScalarQuantizedVectorsFormat(
+            encoding,
+            queryEncoding,
+            Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN,
+            Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH,
+            1,
+            null);
+    super.setUp();
+  }
 
   @Override
   protected Codec getCodec() {
-    return TestUtil.alwaysKnnVectorsFormat(FORMAT);
+    return TestUtil.alwaysKnnVectorsFormat(format);
   }
 
   public void testToString() {