Merge branch 'main' of github.com:elastic/elasticsearch into diskbbq_segment_affinity

Author: tteofili

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

@@ -0,0 +1,145 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", 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", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+package org.elasticsearch.benchmark.vector;
+
+import org.apache.lucene.util.VectorUtil;
+import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
+import org.elasticsearch.common.logging.LogConfigurator;
+import org.elasticsearch.simdvec.ESVectorUtil;
+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.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
+
+import java.io.IOException;
+import java.util.Random;
+import java.util.concurrent.TimeUnit;
+
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+@State(Scope.Benchmark)
+// first iteration is complete garbage, so make sure we really warmup
+@Warmup(iterations = 4, time = 1)
+// real iterations. not useful to spend tons of time here, better to fork more
+@Measurement(iterations = 5, time = 1)
+// engage some noise reduction
+@Fork(value = 1)
+public class DistanceBulkBenchmark {
+
+    static {
+        LogConfigurator.configureESLogging(); // native access requires logging to be initialized
+    }
+
+    @Param({ "384", "782", "1024" })
+    int dims;
+
+    int length;
+
+    int numVectors = 4 * 100;
+    int numQueries = 10;
+
+    float[][] vectors;
+    float[][] queries;
+    float[] distances = new float[4];
+
+    @Setup
+    public void setup() throws IOException {
+        Random random = new Random(123);
+
+        this.length = OptimizedScalarQuantizer.discretize(dims, 64) / 8;
+
+        vectors = new float[numVectors][dims];
+        for (float[] vector : vectors) {
+            for (int i = 0; i < dims; i++) {
+                vector[i] = random.nextFloat();
+            }
+        }
+
+        queries = new float[numQueries][dims];
+        for (float[] query : queries) {
+            for (int i = 0; i < dims; i++) {
+                query[i] = random.nextFloat();
+            }
+        }
+    }
+
+    @Benchmark
+    @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
+    public void squareDistance(Blackhole bh) {
+        for (int j = 0; j < numQueries; j++) {
+            float[] query = queries[j];
+            for (int i = 0; i < numVectors; i++) {
+                float[] vector = vectors[i];
+                float distance = VectorUtil.squareDistance(query, vector);
+                bh.consume(distance);
+            }
+        }
+    }
+
+    @Benchmark
+    @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
+    public void soarDistance(Blackhole bh) {
+        for (int j = 0; j < numQueries; j++) {
+            float[] query = queries[j];
+            for (int i = 0; i < numVectors; i++) {
+                float[] vector = vectors[i];
+                float distance = ESVectorUtil.soarDistance(query, vector, vector, 1.0f, 1.0f);
+                bh.consume(distance);
+            }
+        }
+    }
+
+    @Benchmark
+    @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
+    public void squareDistanceBulk(Blackhole bh) {
+        for (int j = 0; j < numQueries; j++) {
+            float[] query = queries[j];
+            for (int i = 0; i < numVectors; i += 4) {
+                ESVectorUtil.squareDistanceBulk(query, vectors[i], vectors[i + 1], vectors[i + 2], vectors[i + 3], distances);
+                for (float distance : distances) {
+                    bh.consume(distance);
+                }
+
+            }
+        }
+    }
+
+    @Benchmark
+    @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
+    public void soarDistanceBulk(Blackhole bh) {
+        for (int j = 0; j < numQueries; j++) {
+            float[] query = queries[j];
+            for (int i = 0; i < numVectors; i += 4) {
+                ESVectorUtil.soarDistanceBulk(
+                    query,
+                    vectors[i],
+                    vectors[i + 1],
+                    vectors[i + 2],
+                    vectors[i + 3],
+                    vectors[i],
+                    1.0f,
+                    1.0f,
+                    distances
+                );
+                for (float distance : distances) {
+                    bh.consume(distance);
+                }
+
+            }
+        }
+    }
+}

docs/changelog/132167.yaml

@@ -0,0 +1,6 @@
+pr: 132167
+summary: Deal with internally created IN in a different way for EQL
+area: EQL
+type: bug
+issues:
+ - 118621

libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESVectorUtil.java

@@ -293,4 +293,79 @@ public static int quantizeVectorWithIntervals(float[] vector, int[] destination,
         }
         return IMPL.quantizeVectorWithIntervals(vector, destination, lowInterval, upperInterval, bit);
     }
+
+    /**
+     * Bulk computation of square distances between a query vector and four vectors.Result is stored in the provided distances array.
+     *
+     * @param q the query vector
+     * @param v0 the first vector
+     * @param v1 the second vector
+     * @param v2 the third vector
+     * @param v3 the fourth vector
+     * @param distances an array to store the computed square distances, must have length 4
+     *
+     * @throws IllegalArgumentException if the dimensions of the vectors do not match or if the distances array does not have length 4
+     */
+    public static void squareDistanceBulk(float[] q, float[] v0, float[] v1, float[] v2, float[] v3, float[] distances) {
+        if (q.length != v0.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + q.length + "!=" + v0.length);
+        }
+        if (q.length != v1.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + q.length + "!=" + v1.length);
+        }
+        if (q.length != v2.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + q.length + "!=" + v2.length);
+        }
+        if (q.length != v3.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + q.length + "!=" + v3.length);
+        }
+        if (distances.length != 4) {
+            throw new IllegalArgumentException("distances array must have length 4, but was: " + distances.length);
+        }
+        IMPL.squareDistanceBulk(q, v0, v1, v2, v3, distances);
+    }
+
+    /**
+     * Bulk computation of the soar distance for a vector to four centroids
+     * @param v1 the vector
+     * @param c0 the first centroid
+     * @param c1 the second centroid
+     * @param c2 the third centroid
+     * @param c3 the fourth centroid
+     * @param originalResidual the residual with the actually nearest centroid
+     * @param soarLambda the lambda parameter
+     * @param rnorm distance to the nearest centroid
+     * @param distances an array to store the computed soar distances, must have length 4
+     */
+    public static void soarDistanceBulk(
+        float[] v1,
+        float[] c0,
+        float[] c1,
+        float[] c2,
+        float[] c3,
+        float[] originalResidual,
+        float soarLambda,
+        float rnorm,
+        float[] distances
+    ) {
+        if (v1.length != c0.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + v1.length + "!=" + c0.length);
+        }
+        if (v1.length != c1.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + v1.length + "!=" + c1.length);
+        }
+        if (v1.length != c2.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + v1.length + "!=" + c2.length);
+        }
+        if (v1.length != c3.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + v1.length + "!=" + c3.length);
+        }
+        if (v1.length != originalResidual.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + v1.length + "!=" + originalResidual.length);
+        }
+        if (distances.length != 4) {
+            throw new IllegalArgumentException("distances array must have length 4, but was: " + distances.length);
+        }
+        IMPL.soarDistanceBulk(v1, c0, c1, c2, c3, originalResidual, soarLambda, rnorm, distances);
+    }
 }

libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/DefaultESVectorUtilSupport.java

@@ -293,4 +293,30 @@ public int quantizeVectorWithIntervals(float[] vector, int[] destination, float
         }
         return sumQuery;
     }
+
+    @Override
+    public void squareDistanceBulk(float[] query, float[] v0, float[] v1, float[] v2, float[] v3, float[] distances) {
+        distances[0] = VectorUtil.squareDistance(query, v0);
+        distances[1] = VectorUtil.squareDistance(query, v1);
+        distances[2] = VectorUtil.squareDistance(query, v2);
+        distances[3] = VectorUtil.squareDistance(query, v3);
+    }
+
+    @Override
+    public void soarDistanceBulk(
+        float[] v1,
+        float[] c0,
+        float[] c1,
+        float[] c2,
+        float[] c3,
+        float[] originalResidual,
+        float soarLambda,
+        float rnorm,
+        float[] distances
+    ) {
+        distances[0] = soarDistance(v1, c0, originalResidual, soarLambda, rnorm);
+        distances[1] = soarDistance(v1, c1, originalResidual, soarLambda, rnorm);
+        distances[2] = soarDistance(v1, c2, originalResidual, soarLambda, rnorm);
+        distances[3] = soarDistance(v1, c3, originalResidual, soarLambda, rnorm);
+    }
 }

libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorUtilSupport.java

@@ -50,4 +50,17 @@ float calculateOSQLoss(
 
     int quantizeVectorWithIntervals(float[] vector, int[] quantize, float lowInterval, float upperInterval, byte bit);
 
+    void squareDistanceBulk(float[] query, float[] v0, float[] v1, float[] v2, float[] v3, float[] distances);
+
+    void soarDistanceBulk(
+        float[] v1,
+        float[] c0,
+        float[] c1,
+        float[] c2,
+        float[] c3,
+        float[] originalResidual,
+        float soarLambda,
+        float rnorm,
+        float[] distances
+    );
 }