Merge branch 'main' into poc_runtime_function_evaluators_code_only

Author: luigidellaquila

benchmarks/build.gradle

@@ -53,6 +53,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());
                 }