Merge branch 'main' of https://github.com/Jan-Kazlouski-elastic/elasticsearch into feature/hugging-face-chat-completion-integration

Jan-Kazlouski-elastic · Jan-Kazlouski-elastic · commit c5988ed4b0f1 · 2025-05-14T08:55:37.000+03:00
diff --git a/benchmarks/src/main/java/org/elasticsearch/benchmark/vector/OSQScorerBenchmark.java b/benchmarks/src/main/java/org/elasticsearch/benchmark/vector/OSQScorerBenchmark.java
@@ -17,7 +17,7 @@
 import org.apache.lucene.util.VectorUtil;
 import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
 import org.elasticsearch.common.logging.LogConfigurator;
-import org.elasticsearch.simdvec.internal.vectorization.ES91OSQVectorsScorer;
+import org.elasticsearch.simdvec.ES91OSQVectorsScorer;
 import org.elasticsearch.simdvec.internal.vectorization.ESVectorizationProvider;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
diff --git a/build-conventions/settings.gradle b/build-conventions/settings.gradle
@@ -8,7 +8,7 @@
  */
 
 plugins {
-  id "com.gradle.develocity" version "3.19.2"
+  id "com.gradle.develocity" version "4.0.1"
 }
 
 rootProject.name = 'build-conventions'
diff --git a/build-tools-internal/settings.gradle b/build-tools-internal/settings.gradle
@@ -9,7 +9,7 @@ pluginManagement {
 }
 
 plugins {
-  id "com.gradle.develocity" version "3.19.2"
+  id "com.gradle.develocity" version "4.0.1"
 }
 
 dependencyResolutionManagement {
diff --git a/build-tools/settings.gradle b/build-tools/settings.gradle
@@ -10,7 +10,7 @@ pluginManagement {
     includeBuild "../build-conventions"
 }
 plugins {
-  id "com.gradle.develocity" version "3.19.2"
+  id "com.gradle.develocity" version "4.0.1"
 }
 include 'reaper'
 
diff --git a/gradle/build.versions.toml b/gradle/build.versions.toml
@@ -16,7 +16,7 @@ checkstyle = "com.puppycrawl.tools:checkstyle:10.3"
 commmons-io = "commons-io:commons-io:2.2"
 docker-compose = "com.avast.gradle:gradle-docker-compose-plugin:0.17.5"
 forbiddenApis = "de.thetaphi:forbiddenapis:3.8"
-gradle-enterprise = "com.gradle:develocity-gradle-plugin:3.19.2"
+gradle-enterprise = "com.gradle:develocity-gradle-plugin:4.0.1"
 hamcrest = "org.hamcrest:hamcrest:3.0"
 httpcore5 = "org.apache.httpcomponents.core5:httpcore5:5.3.3"
 httpclient5 = "org.apache.httpcomponents.client5:httpclient5:5.4.2"
diff --git a/gradle/verification-metadata.xml b/gradle/verification-metadata.xml
@@ -866,9 +866,11 @@
             <sha256 value="48234cd74e35d91a31a683820a35b5b6d11b55527f32a5b162c6757408b95d7a" origin="Generated by Gradle"/>
          </artifact>
       </component>
-      <component group="com.gradle" name="develocity-gradle-plugin" version="3.19.2">
-         <artifact name="develocity-gradle-plugin-3.19.2.jar">
-            <sha256 value="fc611e8858a901619831ea1d9efd289db7fea28a0ca3d091abdd363365d9512b" origin="Generated by Gradle"/>
+      <component group="com.gradle" name="develocity-gradle-plugin" version="4.0.1">
+         <artifact name="develocity-gradle-plugin-4.0.1.jar">
+            <sha256 value="fc611e8858a901619831ea1d9efd289db7fea28a0ca3d091abdd363365d9512b" origin="Generated by Gradle">
+               <also-trust value="62a98a9df9fbf8c16772787aa5a0009cedcf45c366a0969274bc3b3644c1de25"/>
+            </sha256>
          </artifact>
       </component>
       <component group="com.gradleup.shadow" name="shadow-gradle-plugin" version="8.3.5">
diff --git a/libs/simdvec/src/main/java/org/elasticsearch/simdvec/ES91OSQVectorsScorer.java b/libs/simdvec/src/main/java/org/elasticsearch/simdvec/ES91OSQVectorsScorer.java
@@ -6,7 +6,7 @@
  * 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.simdvec.internal.vectorization;
+package org.elasticsearch.simdvec;
 
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.store.IndexInput;
diff --git a/libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESVectorUtil.java b/libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESVectorUtil.java
@@ -9,11 +9,13 @@
 
 package org.elasticsearch.simdvec;
 
+import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.util.BitUtil;
 import org.apache.lucene.util.Constants;
 import org.elasticsearch.simdvec.internal.vectorization.ESVectorUtilSupport;
 import org.elasticsearch.simdvec.internal.vectorization.ESVectorizationProvider;
 
+import java.io.IOException;
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.MethodType;
@@ -41,6 +43,10 @@ public class ESVectorUtil {
 
     private static final ESVectorUtilSupport IMPL = ESVectorizationProvider.getInstance().getVectorUtilSupport();
 
+    public static ES91OSQVectorsScorer getES91OSQVectorsScorer(IndexInput input, int dimension) throws IOException {
+        return ESVectorizationProvider.getInstance().newES91OSQVectorsScorer(input, dimension);
+    }
+
     public static long ipByteBinByte(byte[] q, byte[] d) {
         if (q.length != d.length * B_QUERY) {
             throw new IllegalArgumentException("vector dimensions incompatible: " + q.length + "!= " + B_QUERY + " x " + d.length);
@@ -211,4 +217,40 @@ public static void centerAndCalculateOSQStatsDp(float[] target, float[] centroid
         assert stats.length == 6;
         IMPL.centerAndCalculateOSQStatsDp(target, centroid, centered, stats);
     }
+
+    /**
+     * Calculates the difference between two vectors and stores the result in a third vector.
+     * @param v1 the first vector
+     * @param v2 the second vector
+     * @param result the result vector, must be the same length as the input vectors
+     */
+    public static void subtract(float[] v1, float[] v2, float[] result) {
+        if (v1.length != v2.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + v1.length + "!=" + v2.length);
+        }
+        if (result.length != v1.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + result.length + "!=" + v1.length);
+        }
+        for (int i = 0; i < v1.length; i++) {
+            result[i] = v1[i] - v2[i];
+        }
+    }
+
+    /**
+     * calculates the spill-over score for a vector and a centroid, given its residual with
+     * its actually nearest centroid
+     * @param v1 the vector
+     * @param centroid the centroid
+     * @param originalResidual the residual with the actually nearest centroid
+     * @return the spill-over score (soar)
+     */
+    public static float soarResidual(float[] v1, float[] centroid, float[] originalResidual) {
+        if (v1.length != centroid.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + v1.length + "!=" + centroid.length);
+        }
+        if (originalResidual.length != v1.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + originalResidual.length + "!=" + v1.length);
+        }
+        return IMPL.soarResidual(v1, centroid, originalResidual);
+    }
 }
diff --git a/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/DefaultESVectorUtilSupport.java b/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/DefaultESVectorUtilSupport.java
@@ -138,6 +138,18 @@ public void centerAndCalculateOSQStatsDp(float[] target, float[] centroid, float
         stats[5] = centroidDot;
     }
 
+    @Override
+    public float soarResidual(float[] v1, float[] centroid, float[] originalResidual) {
+        assert v1.length == centroid.length;
+        assert v1.length == originalResidual.length;
+        float proj = 0;
+        for (int i = 0; i < v1.length; i++) {
+            float djk = v1[i] - centroid[i];
+            proj = fma(djk, originalResidual[i], proj);
+        }
+        return proj;
+    }
+
     public static int ipByteBitImpl(byte[] q, byte[] d) {
         return ipByteBitImpl(q, d, 0);
     }
diff --git a/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/DefaultESVectorizationProvider.java b/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/DefaultESVectorizationProvider.java
@@ -10,6 +10,7 @@
 package org.elasticsearch.simdvec.internal.vectorization;
 
 import org.apache.lucene.store.IndexInput;
+import org.elasticsearch.simdvec.ES91OSQVectorsScorer;
 
 import java.io.IOException;
 
diff --git a/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorUtilSupport.java b/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorUtilSupport.java
@@ -28,4 +28,7 @@ public interface ESVectorUtilSupport {
     void centerAndCalculateOSQStatsEuclidean(float[] target, float[] centroid, float[] centered, float[] stats);
 
     void centerAndCalculateOSQStatsDp(float[] target, float[] centroid, float[] centered, float[] stats);
+
+    float soarResidual(float[] v1, float[] centroid, float[] originalResidual);
+
 }
diff --git a/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorizationProvider.java b/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorizationProvider.java
@@ -10,6 +10,7 @@
 package org.elasticsearch.simdvec.internal.vectorization;
 
 import org.apache.lucene.store.IndexInput;
+import org.elasticsearch.simdvec.ES91OSQVectorsScorer;
 
 import java.io.IOException;
 import java.util.Objects;
diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorizationProvider.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorizationProvider.java
@@ -13,6 +13,7 @@
 import org.apache.lucene.util.Constants;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
+import org.elasticsearch.simdvec.ES91OSQVectorsScorer;
 
 import java.io.IOException;
 import java.util.Locale;
diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentES91OSQVectorsScorer.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentES91OSQVectorsScorer.java
@@ -20,6 +20,7 @@
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.util.VectorUtil;
 import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
+import org.elasticsearch.simdvec.ES91OSQVectorsScorer;
 
 import java.io.IOException;
 import java.lang.foreign.MemorySegment;
diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/PanamaESVectorUtilSupport.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/PanamaESVectorUtilSupport.java
@@ -367,6 +367,49 @@ public float calculateOSQLoss(float[] target, float[] interval, float step, floa
         return (1f - lambda) * xe * xe / norm2 + lambda * e;
     }
 
+    @Override
+    public float soarResidual(float[] v1, float[] centroid, float[] originalResidual) {
+        assert v1.length == centroid.length;
+        assert v1.length == originalResidual.length;
+        float proj = 0;
+        int i = 0;
+        if (v1.length > 2 * FLOAT_SPECIES.length()) {
+            FloatVector projVec1 = FloatVector.zero(FLOAT_SPECIES);
+            FloatVector projVec2 = FloatVector.zero(FLOAT_SPECIES);
+            int unrolledLimit = FLOAT_SPECIES.loopBound(v1.length) - FLOAT_SPECIES.length();
+            for (; i < unrolledLimit; i += 2 * FLOAT_SPECIES.length()) {
+                // one
+                FloatVector v1Vec0 = FloatVector.fromArray(FLOAT_SPECIES, v1, i);
+                FloatVector centroidVec0 = FloatVector.fromArray(FLOAT_SPECIES, centroid, i);
+                FloatVector originalResidualVec0 = FloatVector.fromArray(FLOAT_SPECIES, originalResidual, i);
+                FloatVector djkVec0 = v1Vec0.sub(centroidVec0);
+                projVec1 = fma(djkVec0, originalResidualVec0, projVec1);
+
+                // two
+                FloatVector v1Vec1 = FloatVector.fromArray(FLOAT_SPECIES, v1, i + FLOAT_SPECIES.length());
+                FloatVector centroidVec1 = FloatVector.fromArray(FLOAT_SPECIES, centroid, i + FLOAT_SPECIES.length());
+                FloatVector originalResidualVec1 = FloatVector.fromArray(FLOAT_SPECIES, originalResidual, i + FLOAT_SPECIES.length());
+                FloatVector djkVec1 = v1Vec1.sub(centroidVec1);
+                projVec2 = fma(djkVec1, originalResidualVec1, projVec2);
+            }
+            // vector tail
+            for (; i < FLOAT_SPECIES.loopBound(v1.length); i += FLOAT_SPECIES.length()) {
+                FloatVector v1Vec = FloatVector.fromArray(FLOAT_SPECIES, v1, i);
+                FloatVector centroidVec = FloatVector.fromArray(FLOAT_SPECIES, centroid, i);
+                FloatVector originalResidualVec = FloatVector.fromArray(FLOAT_SPECIES, originalResidual, i);
+                FloatVector djkVec = v1Vec.sub(centroidVec);
+                projVec1 = fma(djkVec, originalResidualVec, projVec1);
+            }
+            proj += projVec1.add(projVec2).reduceLanes(ADD);
+        }
+        // tail
+        for (; i < v1.length; i++) {
+            float djk = v1[i] - centroid[i];
+            proj = fma(djk, originalResidual[i], proj);
+        }
+        return proj;
+    }
+
     private static final VectorSpecies<Byte> BYTE_SPECIES_128 = ByteVector.SPECIES_128;
     private static final VectorSpecies<Byte> BYTE_SPECIES_256 = ByteVector.SPECIES_256;
 
diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/PanamaESVectorizationProvider.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/PanamaESVectorizationProvider.java
@@ -11,6 +11,7 @@
 
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.store.MemorySegmentAccessInput;
+import org.elasticsearch.simdvec.ES91OSQVectorsScorer;
 
 import java.io.IOException;
 import java.lang.foreign.MemorySegment;
diff --git a/libs/simdvec/src/test/java/org/elasticsearch/simdvec/ESVectorUtilTests.java b/libs/simdvec/src/test/java/org/elasticsearch/simdvec/ESVectorUtilTests.java
@@ -268,6 +268,22 @@ public void testOsqGridPoints() {
         }
     }
 
+    public void testSoarOverspillScore() {
+        int size = random().nextInt(128, 512);
+        float deltaEps = 1e-5f * size;
+        var vector = new float[size];
+        var centroid = new float[size];
+        var preResidual = new float[size];
+        for (int i = 0; i < size; ++i) {
+            vector[i] = random().nextFloat();
+            centroid[i] = random().nextFloat();
+            preResidual[i] = random().nextFloat();
+        }
+        var expected = defaultedProvider.getVectorUtilSupport().soarResidual(vector, centroid, preResidual);
+        var result = defOrPanamaProvider.getVectorUtilSupport().soarResidual(vector, centroid, preResidual);
+        assertEquals(expected, result, deltaEps);
+    }
+
     void testIpByteBinImpl(ToLongBiFunction<byte[], byte[]> ipByteBinFunc) {
         int iterations = atLeast(50);
         for (int i = 0; i < iterations; i++) {
diff --git a/libs/simdvec/src/test/java/org/elasticsearch/simdvec/internal/vectorization/ES91OSQVectorScorerTests.java b/libs/simdvec/src/test/java/org/elasticsearch/simdvec/internal/vectorization/ES91OSQVectorScorerTests.java
@@ -16,6 +16,7 @@
 import org.apache.lucene.store.IndexOutput;
 import org.apache.lucene.store.MMapDirectory;
 import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
+import org.elasticsearch.simdvec.ES91OSQVectorsScorer;
 
 import static org.hamcrest.Matchers.lessThan;
 
diff --git a/plugins/examples/settings.gradle b/plugins/examples/settings.gradle
@@ -8,7 +8,7 @@
  */
 
 plugins {
-  id "com.gradle.develocity" version "3.19.2"
+  id "com.gradle.develocity" version "4.0.1"
 }
 
 // Include all subdirectories as example projects
diff --git a/server/src/main/java/module-info.java b/server/src/main/java/module-info.java
@@ -454,7 +454,8 @@
             org.elasticsearch.index.codec.vectors.es816.ES816BinaryQuantizedVectorsFormat,
             org.elasticsearch.index.codec.vectors.es816.ES816HnswBinaryQuantizedVectorsFormat,
             org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat,
-            org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat;
+            org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat,
+            org.elasticsearch.index.codec.vectors.IVFVectorsFormat;
 
     provides org.apache.lucene.codecs.Codec
         with
diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsReader.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsReader.java
diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsWriter.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsWriter.java
diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsFormat.java
diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsReader.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsReader.java
diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsWriter.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsWriter.java
diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/NeighborQueue.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/NeighborQueue.java
diff --git a/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat b/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat
diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/IVFVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/IVFVectorsFormatTests.java
diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/NeighborQueueTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/NeighborQueueTests.java
diff --git a/settings.gradle b/settings.gradle

Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,7 @@`
`8`	`8`	`*/`
`9`	`9`
`10`	`10`	`plugins {`
`11`		`- id "com.gradle.develocity" version "3.19.2"`
	`11`	`+ id "com.gradle.develocity" version "4.0.1"`
`12`	`12`	`}`
`13`	`13`
`14`	`14`	`rootProject.name = 'build-conventions'`
Original file line number	Diff line number	Diff line change
`@@ -9,7 +9,7 @@ pluginManagement {`
`9`	`9`	`}`
`10`	`10`
`11`	`11`	`plugins {`
`12`		`- id "com.gradle.develocity" version "3.19.2"`
	`12`	`+ id "com.gradle.develocity" version "4.0.1"`
`13`	`13`	`}`
`14`	`14`
`15`	`15`	`dependencyResolutionManagement {`
Original file line number	Diff line number	Diff line change
`@@ -10,7 +10,7 @@ pluginManagement {`
`10`	`10`	`includeBuild "../build-conventions"`
`11`	`11`	`}`
`12`	`12`	`plugins {`
`13`		`- id "com.gradle.develocity" version "3.19.2"`
	`13`	`+ id "com.gradle.develocity" version "4.0.1"`
`14`	`14`	`}`
`15`	`15`	`include 'reaper'`
`16`	`16`