Fix VectorScorerOSQBenchmarkTests (elastic#142742)

This PR fixes several VectorScorerOSQBenchmarkTests failures.

The native C code (score_common.h) is either compiled with -O3 (e.g. on aarch64), which enables FMA (fused multiply-add) instructions, or uses FMA SIMD instructions directly. Java's float arithmetic rounds at every operation per IEEE 754, but FMA computes a * b + c with a single rounding, preserving extra intermediate precision.

The score correction formula involves terms like ax * ay * dims + ay * lx * tcs + .... When correction values come from random bytes, they can represent extreme floats (e.g., ±1e+30). These large terms nearly cancel, and the FMA vs non-FMA rounding difference becomes catastrophic — the scalar gets score ≈ -1 (result 0.0) while the native gets score ≈ 5.5e+37 (result 2.7e+37).

This is NOT an issue in production — real quantized corrections from OptimizedScalarQuantizer produce small, well-behaved floats where FMA differences are negligible (the ESNextOSQVectorsScorerTests.testScoreBulk test validates this with real data).

The PR extracts the same input data generation used in ESNextOSQVectorsScorerTests.testScoreBulk, extracting it and exposing it via test fixture, so both tests share the same input data setup.

Fixes elastic#142289
Fixes elastic#142413
Fixes elastic#142490
Fixes elastic#142491
Fixes elastic#142492
Fixes elastic#142587
Fixes elastic#142588
Fixes elastic#142589

Author: ldematte

benchmarks/build.gradle

@@ -51,6 +51,7 @@ dependencies {
   api(project(':x-pack:plugin:logsdb'))
   implementation project(path: ':libs:native')
   implementation project(path: ':libs:simdvec')
+  implementation (testFixtures(project(path: ':libs:simdvec')))
   implementation project(path: ':libs:swisshash')
   implementation project(path: ':libs:exponential-histogram')
   implementation(project(':x-pack:plugin:searchable-snapshots')) {

benchmarks/src/main/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerOSQBenchmark.java

@@ -16,12 +16,13 @@
 import org.apache.lucene.store.IndexOutput;
 import org.apache.lucene.store.MMapDirectory;
 import org.apache.lucene.store.NIOFSDirectory;
-import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
+import org.apache.lucene.util.VectorUtil;
 import org.elasticsearch.common.logging.LogConfigurator;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat;
 import org.elasticsearch.simdvec.ESNextOSQVectorsScorer;
 import org.elasticsearch.simdvec.internal.vectorization.ESVectorizationProvider;
+import org.elasticsearch.simdvec.internal.vectorization.VectorScorerTestUtils;
 import org.elasticsearch.xpack.searchablesnapshots.store.SearchableSnapshotDirectoryFactory;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
@@ -42,6 +43,11 @@
 import java.util.Random;
 import java.util.concurrent.TimeUnit;
 
+import static org.elasticsearch.simdvec.internal.vectorization.VectorScorerTestUtils.createOSQIndexData;
+import static org.elasticsearch.simdvec.internal.vectorization.VectorScorerTestUtils.createOSQQueryData;
+import static org.elasticsearch.simdvec.internal.vectorization.VectorScorerTestUtils.randomVector;
+import static org.elasticsearch.simdvec.internal.vectorization.VectorScorerTestUtils.writeBulkOSQVectorData;
+
 @BenchmarkMode(Mode.Throughput)
 @OutputTimeUnit(TimeUnit.MILLISECONDS)
 @State(Scope.Benchmark)
@@ -72,7 +78,7 @@ public enum VectorImplementation {
     public int dims;
 
     @Param({ "1", "2", "4" })
-    public int bits;
+    public byte bits;
 
     int bulkSize = ESNextOSQVectorsScorer.BULK_SIZE;
 
@@ -90,9 +96,7 @@ public enum VectorImplementation {
 
     int length;
 
-    byte[][] binaryVectors;
-    byte[][] binaryQueries;
-    OptimizedScalarQuantizer.QuantizationResult result;
+    VectorScorerTestUtils.OSQVectorData[] binaryQueries;
     float centroidDp;
 
     byte[] scratch;
@@ -111,17 +115,12 @@ public void setup() throws IOException {
     }
 
     void setup(Random random) throws IOException {
-        this.length = switch (bits) {
-            case 1 -> ESNextDiskBBQVectorsFormat.QuantEncoding.ONE_BIT_4BIT_QUERY.getDocPackedLength(dims);
-            case 2 -> ESNextDiskBBQVectorsFormat.QuantEncoding.TWO_BIT_4BIT_QUERY.getDocPackedLength(dims);
-            case 4 -> ESNextDiskBBQVectorsFormat.QuantEncoding.FOUR_BIT_SYMMETRIC.getDocPackedLength(dims);
-            default -> throw new IllegalArgumentException("Unsupported bits: " + bits);
-        };
+        this.length = ESNextDiskBBQVectorsFormat.QuantEncoding.fromBits(bits).getDocPackedLength(dims);
 
-        binaryVectors = new byte[numVectors][length];
-        for (byte[] binaryVector : binaryVectors) {
-            random.nextBytes(binaryVector);
-        }
+        final float[] centroid = new float[dims];
+        randomVector(random, centroid, similarityFunction);
+
+        var quantizer = new org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer(similarityFunction);
 
         directory = switch (directoryType) {
             case MMAP -> new MMapDirectory(createTempDirectory("vectorDataMmap"));
@@ -130,35 +129,27 @@ void setup(Random random) throws IOException {
         };
 
         try (IndexOutput output = directory.createOutput("vectors", IOContext.DEFAULT)) {
-            byte[] correctionBytes = new byte[16 * bulkSize];
+            VectorScorerTestUtils.OSQVectorData[] vectors = new VectorScorerTestUtils.OSQVectorData[bulkSize];
             for (int i = 0; i < numVectors; i += bulkSize) {
                 for (int j = 0; j < bulkSize; j++) {
-                    output.writeBytes(binaryVectors[i + j], 0, binaryVectors[i + j].length);
+                    var vector = new float[dims];
+                    randomVector(random, vector, similarityFunction);
+                    vectors[j] = createOSQIndexData(vector, centroid, quantizer, dims, bits, length);
                 }
-                random.nextBytes(correctionBytes);
-                output.writeBytes(correctionBytes, 0, correctionBytes.length);
+                writeBulkOSQVectorData(bulkSize, output, vectors);
             }
             CodecUtil.writeFooter(output);
         }
         input = directory.openInput("vectors", IOContext.DEFAULT);
-        int binaryQueryLength = switch (bits) {
-            case 1 -> ESNextDiskBBQVectorsFormat.QuantEncoding.ONE_BIT_4BIT_QUERY.getQueryPackedLength(dims);
-            case 2 -> ESNextDiskBBQVectorsFormat.QuantEncoding.TWO_BIT_4BIT_QUERY.getQueryPackedLength(dims);
-            case 4 -> ESNextDiskBBQVectorsFormat.QuantEncoding.FOUR_BIT_SYMMETRIC.getQueryPackedLength(dims);
-            default -> throw new IllegalArgumentException("Unsupported bits: " + bits);
-        };
+        int binaryQueryLength = ESNextDiskBBQVectorsFormat.QuantEncoding.fromBits(bits).getQueryPackedLength(dims);
 
-        binaryQueries = new byte[numVectors][binaryQueryLength];
-        for (byte[] binaryQuery : binaryQueries) {
-            random.nextBytes(binaryQuery);
+        binaryQueries = new VectorScorerTestUtils.OSQVectorData[numVectors];
+        var query = new float[dims];
+        for (int i = 0; i < numVectors; ++i) {
+            randomVector(random, query, similarityFunction);
+            binaryQueries[i] = createOSQQueryData(query, centroid, quantizer, dims, (byte) 4, binaryQueryLength);
         }
-        result = new OptimizedScalarQuantizer.QuantizationResult(
-            random.nextFloat(),
-            random.nextFloat(),
-            random.nextFloat(),
-            Short.toUnsignedInt((short) random.nextInt())
-        );
-        centroidDp = random.nextFloat();
+        centroidDp = VectorUtil.dotProduct(centroid, centroid);
 
         scratch = new byte[length];
         final int docBits;
@@ -202,14 +193,14 @@ public float[] score() throws IOException {
         for (int j = 0; j < numQueries; j++) {
             input.seek(0);
             for (int i = 0; i < numVectors; i++) {
-                float qDist = scorer.quantizeScore(binaryQueries[j]);
+                float qDist = scorer.quantizeScore(binaryQueries[j].quantizedVector());
                 input.readFloats(corrections, 0, corrections.length);
                 int addition = Short.toUnsignedInt(input.readShort());
                 float score = scorer.score(
-                    result.lowerInterval(),
-                    result.upperInterval(),
-                    result.quantizedComponentSum(),
-                    result.additionalCorrection(),
+                    binaryQueries[j].lowerInterval(),
+                    binaryQueries[j].upperInterval(),
+                    binaryQueries[j].quantizedComponentSum(),
+                    binaryQueries[j].additionalCorrection(),
                     similarityFunction,
                     centroidDp,
                     corrections[0],
@@ -231,11 +222,11 @@ public float[] bulkScore() throws IOException {
             input.seek(0);
             for (int i = 0; i < numVectors; i += scratchScores.length) {
                 scorer.scoreBulk(
-                    binaryQueries[j],
-                    result.lowerInterval(),
-                    result.upperInterval(),
-                    result.quantizedComponentSum(),
-                    result.additionalCorrection(),
+                    binaryQueries[j].quantizedVector(),
+                    binaryQueries[j].lowerInterval(),
+                    binaryQueries[j].upperInterval(),
+                    binaryQueries[j].quantizedComponentSum(),
+                    binaryQueries[j].additionalCorrection(),
                     similarityFunction,
                     centroidDp,
                     scratchScores

benchmarks/src/test/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerOSQBenchmarkTests.java

@@ -15,6 +15,7 @@
 import org.apache.lucene.util.Constants;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.test.junit.annotations.TestLogging;
 import org.junit.BeforeClass;
 import org.openjdk.jmh.annotations.Param;
 
@@ -24,17 +25,22 @@
 
 import static org.elasticsearch.common.util.CollectionUtils.appendToCopy;
 
+@TestLogging(
+    reason = "Noisy logging",
+    value = "org.elasticsearch.env.NodeEnvironment:WARN,org.elasticsearch.xpack.searchablesnapshots.cache.full.PersistentCache:WARN"
+)
 public class VectorScorerOSQBenchmarkTests extends ESTestCase {
 
+    private static final int REPETITIONS = 10;
     private final float deltaPercent = 0.1f;
     private final int dims;
-    private final int bits;
+    private final byte bits;
     private final VectorScorerOSQBenchmark.DirectoryType directoryType;
     private final VectorSimilarityFunction similarityFunction;
 
     public VectorScorerOSQBenchmarkTests(
         int dims,
-        int bits,
+        byte bits,
         VectorScorerOSQBenchmark.DirectoryType directoryType,
         VectorSimilarityFunction similarityFunction
     ) {
@@ -50,7 +56,7 @@ public static void skipWindows() {
     }
 
     public void testSingleScalarVsVectorized() throws Exception {
-        for (int i = 0; i < 100; i++) {
+        for (int i = 0; i < REPETITIONS; i++) {
             var seed = randomLong();
 
             var scalar = new VectorScorerOSQBenchmark();
@@ -85,7 +91,7 @@ public void testSingleScalarVsVectorized() throws Exception {
     }
 
     public void testBulkScalarVsVectorized() throws Exception {
-        for (int i = 0; i < 100; i++) {
+        for (int i = 0; i < REPETITIONS; i++) {
             var seed = randomLong();
 
             var scalar = new VectorScorerOSQBenchmark();
@@ -128,7 +134,7 @@ public static Iterable<Object[]> parametersFactory() {
 
             return () -> Arrays.stream(dims)
                 .map(Integer::parseInt)
-                .flatMap(d -> Arrays.stream(bits).map(Integer::parseInt).map(b -> List.<Object>of(d, b)))
+                .flatMap(d -> Arrays.stream(bits).map(Byte::parseByte).map(b -> List.<Object>of(d, b)))
                 .flatMap(params -> Arrays.stream(VectorScorerOSQBenchmark.DirectoryType.values()).map(dir -> appendToCopy(params, dir)))
                 .flatMap(params -> Arrays.stream(VectorSimilarityFunction.values()).map(f -> appendToCopy(params, f).toArray()))
                 .iterator();

libs/simdvec/build.gradle

@@ -12,16 +12,20 @@ import org.elasticsearch.gradle.internal.precommit.CheckForbiddenApisTask
 apply plugin: 'elasticsearch.publish'
 apply plugin: 'elasticsearch.build'
 apply plugin: 'elasticsearch.mrjar'
+apply plugin: 'java-test-fixtures'
 
 dependencies {
   implementation project(':libs:native')
   implementation project(':libs:logging')
   implementation "org.apache.lucene:lucene-core:${versions.lucene}"
 
+  testImplementation(testArtifact(project(':x-pack:plugin:searchable-snapshots')))
   testImplementation(project(":test:framework")) {
     exclude group: 'org.elasticsearch', module: 'native'
   }
-  testImplementation(testArtifact(project(':x-pack:plugin:searchable-snapshots')))
+  testFixturesImplementation(project(":test:framework")) {
+    exclude group: 'org.elasticsearch', module: 'native'
+  }
 }
 
 // compileMain21Java does not exist within idea (see MrJarPlugin) so we cannot reference directly by name

libs/simdvec/src/test/java/org/elasticsearch/simdvec/internal/vectorization/ES91OSQVectorScorerTests.java

@@ -120,7 +120,7 @@ public void testScore() throws Exception {
                     long qDist = defaultScorer.quantizeScore(quantizeQuery);
                     slice.readFloats(floatScratch, 0, 3);
                     int quantizedComponentSum = slice.readShort();
-                    float defaulScore = defaultScorer.score(
+                    float defaultScore = defaultScorer.score(
                         queryCorrections.lowerInterval(),
                         queryCorrections.upperInterval(),
                         queryCorrections.quantizedComponentSum(),
@@ -149,7 +149,7 @@ public void testScore() throws Exception {
                         floatScratch[2],
                         qDist
                     );
-                    assertEquals(defaulScore, panamaScore, 1e-2f);
+                    assertEquals(defaultScore, panamaScore, 1e-2f);
                     assertEquals(((long) (i + 1) * (length + 14)), slice.getFilePointer());
                     assertEquals(padding + ((long) (i + 1) * (length + 14)), in.getFilePointer());
                 }
@@ -234,9 +234,7 @@ public void testScoreBulk() throws Exception {
                         scoresPanama
                     );
                     assertEquals(defaultMaxScore, panamaMaxScore, 1e-2f);
-                    for (int j = 0; j < bulkSize; j++) {
-                        assertEquals(scoresDefault[j], scoresPanama[j], 1e-2f);
-                    }
+                    assertArrayEqualsPercent(scoresDefault, scoresPanama, 0.05f, 1e-2f);
                     assertEquals(((long) (bulkSize) * (length + 14)), slice.getFilePointer());
                     assertEquals(padding + ((long) (i + bulkSize) * (length + 14)), in.getFilePointer());
                 }