Merge branch 'main' into lru

Author: jdconrad

benchmarks/build.gradle

@@ -12,6 +12,8 @@ apply plugin: org.elasticsearch.gradle.internal.ElasticsearchJavaBasePlugin
 apply plugin: 'java-library'
 apply plugin: 'application'
 
+var os = org.gradle.internal.os.OperatingSystem.current()
+
 application {
   mainClass = 'org.openjdk.jmh.Main'
 }
@@ -39,6 +41,7 @@ dependencies {
   api(project(':x-pack:plugin:ql'))
   api(project(':x-pack:plugin:esql'))
   api(project(':x-pack:plugin:esql:compute'))
+  implementation project(path: ':libs:elasticsearch-vec')
   expression(project(path: ':modules:lang-expression', configuration: 'zip'))
   painless(project(path: ':modules:lang-painless', configuration: 'zip'))
   api "org.openjdk.jmh:jmh-core:$versions.jmh"
@@ -73,6 +76,16 @@ tasks.named("run").configure {
   executable = "${BuildParams.runtimeJavaHome}/bin/java"
   args << "-Dplugins.dir=${buildDir}/plugins" << "-Dtests.index=${buildDir}/index"
   dependsOn "copyExpression", "copyPainless"
+  systemProperty 'java.library.path', file("../libs/native/libraries/build/platform/${platformName()}-${os.arch}")
+}
+
+String platformName() {
+  String name = System.getProperty("os.name");
+  if (name.startsWith("Mac")) {
+    return "darwin";
+  } else {
+    return name.toLowerCase(Locale.ROOT);
+  }
 }
 
 spotless {

benchmarks/src/main/java/org/elasticsearch/benchmark/vector/VectorScorerBenchmark.java

@@ -0,0 +1,188 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.benchmark.vector;
+
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.store.IndexOutput;
+import org.apache.lucene.store.MMapDirectory;
+import org.apache.lucene.util.quantization.ScalarQuantizedVectorSimilarity;
+import org.elasticsearch.common.logging.LogConfigurator;
+import org.elasticsearch.core.IOUtils;
+import org.elasticsearch.vec.VectorScorer;
+import org.elasticsearch.vec.VectorScorerFactory;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Warmup;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.TimeUnit;
+
+import static org.elasticsearch.vec.VectorSimilarityType.DOT_PRODUCT;
+import static org.elasticsearch.vec.VectorSimilarityType.EUCLIDEAN;
+
+@Fork(value = 1, jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
+@Warmup(iterations = 3, time = 3)
+@Measurement(iterations = 5, time = 3)
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(TimeUnit.MICROSECONDS)
+@State(Scope.Thread)
+/**
+ * Benchmark that compares various scalar quantized vector similarity function
+ * implementations;: scalar, lucene's panama-ized, and Elasticsearch's native.
+ * Run with ./gradlew -p benchmarks run --args 'VectorScorerBenchmark'
+ */
+public class VectorScorerBenchmark {
+
+    static {
+        LogConfigurator.configureESLogging(); // native access requires logging to be initialized
+    }
+
+    @Param({ "96", "768", "1024" })
+    int dims;
+    int size = 2; // there are only two vectors to compare
+
+    Directory dir;
+    IndexInput in;
+    VectorScorerFactory factory;
+
+    byte[] vec1;
+    byte[] vec2;
+    float vec1Offset;
+    float vec2Offset;
+    float scoreCorrectionConstant;
+
+    ScalarQuantizedVectorSimilarity luceneDotScorer;
+    ScalarQuantizedVectorSimilarity luceneSqrScorer;
+    VectorScorer nativeDotScorer;
+    VectorScorer nativeSqrScorer;
+
+    @Setup
+    public void setup() throws IOException {
+        var optionalVectorScorerFactory = VectorScorerFactory.instance();
+        if (optionalVectorScorerFactory.isEmpty()) {
+            String msg = "JDK=["
+                + Runtime.version()
+                + "], os.name=["
+                + System.getProperty("os.name")
+                + "], os.arch=["
+                + System.getProperty("os.arch")
+                + "]";
+            throw new AssertionError("Vector scorer factory not present. Cannot run the benchmark. " + msg);
+        }
+        factory = optionalVectorScorerFactory.get();
+        scoreCorrectionConstant = 1f;
+        vec1 = new byte[dims];
+        vec2 = new byte[dims];
+
+        ThreadLocalRandom.current().nextBytes(vec1);
+        ThreadLocalRandom.current().nextBytes(vec2);
+        vec1Offset = ThreadLocalRandom.current().nextFloat();
+        vec2Offset = ThreadLocalRandom.current().nextFloat();
+
+        dir = new MMapDirectory(Files.createTempDirectory("nativeScalarQuantBench"));
+        try (IndexOutput out = dir.createOutput("vector.data", IOContext.DEFAULT)) {
+            out.writeBytes(vec1, 0, vec1.length);
+            out.writeInt(Float.floatToIntBits(vec1Offset));
+            out.writeBytes(vec2, 0, vec2.length);
+            out.writeInt(Float.floatToIntBits(vec2Offset));
+        }
+        in = dir.openInput("vector.data", IOContext.DEFAULT);
+
+        luceneDotScorer = ScalarQuantizedVectorSimilarity.fromVectorSimilarity(
+            VectorSimilarityFunction.DOT_PRODUCT,
+            scoreCorrectionConstant
+        );
+        luceneSqrScorer = ScalarQuantizedVectorSimilarity.fromVectorSimilarity(VectorSimilarityFunction.EUCLIDEAN, scoreCorrectionConstant);
+        nativeDotScorer = factory.getScalarQuantizedVectorScorer(dims, size, scoreCorrectionConstant, DOT_PRODUCT, in).get();
+        nativeSqrScorer = factory.getScalarQuantizedVectorScorer(dims, size, scoreCorrectionConstant, EUCLIDEAN, in).get();
+
+        // sanity
+        var f1 = dotProductLucene();
+        var f2 = dotProductNative();
+        var f3 = dotProductScalar();
+        if (f1 != f2) {
+            throw new AssertionError("lucene[" + f1 + "] != " + "native[" + f2 + "]");
+        }
+        if (f1 != f3) {
+            throw new AssertionError("lucene[" + f1 + "] != " + "scalar[" + f3 + "]");
+        }
+        // square distance
+        f1 = squareDistanceLucene();
+        f2 = squareDistanceNative();
+        f3 = squareDistanceScalar();
+        if (f1 != f2) {
+            throw new AssertionError("lucene[" + f1 + "] != " + "native[" + f2 + "]");
+        }
+        if (f1 != f3) {
+            throw new AssertionError("lucene[" + f1 + "] != " + "scalar[" + f3 + "]");
+        }
+    }
+
+    @TearDown
+    public void teardown() throws IOException {
+        IOUtils.close(dir, in);
+    }
+
+    @Benchmark
+    public float dotProductLucene() {
+        return luceneDotScorer.score(vec1, vec1Offset, vec2, vec2Offset);
+    }
+
+    @Benchmark
+    public float dotProductNative() throws IOException {
+        return nativeDotScorer.score(0, 1);
+    }
+
+    @Benchmark
+    public float dotProductScalar() {
+        int dotProduct = 0;
+        for (int i = 0; i < vec1.length; i++) {
+            dotProduct += vec1[i] * vec2[i];
+        }
+        float adjustedDistance = dotProduct * scoreCorrectionConstant + vec1Offset + vec2Offset;
+        return (1 + adjustedDistance) / 2;
+    }
+
+    // -- square distance
+
+    @Benchmark
+    public float squareDistanceLucene() {
+        return luceneSqrScorer.score(vec1, vec1Offset, vec2, vec2Offset);
+    }
+
+    @Benchmark
+    public float squareDistanceNative() throws IOException {
+        return nativeSqrScorer.score(0, 1);
+    }
+
+    @Benchmark
+    public float squareDistanceScalar() {
+        int squareDistance = 0;
+        for (int i = 0; i < vec1.length; i++) {
+            int diff = vec1[i] - vec2[i];
+            squareDistance += diff * diff;
+        }
+        float adjustedDistance = squareDistance * scoreCorrectionConstant;
+        return 1 / (1f + adjustedDistance);
+    }
+}

build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/InternalDistributionModuleCheckTaskProvider.java

@@ -63,6 +63,7 @@ public class InternalDistributionModuleCheckTaskProvider {
         "org.elasticsearch.securesm",
         "org.elasticsearch.server",
         "org.elasticsearch.tdigest",
+        "org.elasticsearch.vec",
         "org.elasticsearch.xcontent"
     );
 

build-tools-internal/src/test/groovy/org/elasticsearch/gradle/internal/doc/DocSnippetTaskSpec.groovy

@@ -529,13 +529,15 @@ GET /_analyze
     {
       "type": "mapping",
       "mappings": [
-        "٠ => 0",
-        "١ => 1",
-        "٢ => 2"
+        "e => 0",
+        "m => 1",
+        "p => 2",
+        "t => 3",
+        "y => 4"
       ]
     }
   ],
-  "text": "My license plate is ٢٥٠١٥"
+  "text": "My license plate is empty"
 }
 ----
 """
@@ -551,13 +553,15 @@ GET /_analyze
     {
       "type": "mapping",
       "mappings": [
-        "٠ => 0",
-        "١ => 1",
-        "٢ => 2"
+        "e => 0",
+        "m => 1",
+        "p => 2",
+        "t => 3",
+        "y => 4"
       ]
     }
   ],
-  "text": "My license plate is ٢٥٠١٥"
+  "text": "My license plate is empty"
 }"""
     }
 

build-tools-internal/src/test/groovy/org/elasticsearch/gradle/internal/doc/RestTestsFromDocSnippetTaskSpec.groovy

@@ -454,11 +454,11 @@ Requires the `manage_ml` cluster privilege. This privilege is included in the
 
 `<job_id>`::
 (Required, string)
-include::{es-repo-dir}/ml/ml-shared.asciidoc[tag=job-id-anomaly-detection]
+include::{es-ref-dir}/ml/ml-shared.asciidoc[tag=job-id-anomaly-detection]
 
 `<snapshot_id>`::
 (Required, string)
-include::{es-repo-dir}/ml/ml-shared.asciidoc[tag=snapshot-id]
+include::{es-ref-dir}/ml/ml-shared.asciidoc[tag=snapshot-id]
 
 [[ml-update-snapshot-request-body]]
 == {api-request-body-title}
@@ -470,7 +470,7 @@ The following properties can be updated after the model snapshot is created:
 
 `retain`::
 (Optional, Boolean)
-include::{es-repo-dir}/ml/ml-shared.asciidoc[tag=retain]
+include::{es-ref-dir}/ml/ml-shared.asciidoc[tag=retain]
 
 
 [[ml-update-snapshot-example]]