Merge remote-tracking branch 'origin/main' into detect-mp-in-rest-tests

Author: ywangd

docs/changelog/122381.yaml

@@ -0,0 +1,6 @@
+pr: 122381
+summary: Adds implementations of dotProduct and cosineSimilarity painless methods to operate on float vectors for byte fields
+area: Vector Search
+type: enhancement
+issues:
+ - 117274

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

@@ -80,6 +80,19 @@ public static float ipFloatBit(float[] q, byte[] d) {
         return IMPL.ipFloatBit(q, d);
     }
 
+    /**
+     * Compute the inner product of two vectors, where the query vector is a float vector and the document vector is a byte vector.
+     * @param q the query vector
+     * @param d the document vector
+     * @return the inner product of the two vectors
+     */
+    public static float ipFloatByte(float[] q, byte[] d) {
+        if (q.length != d.length) {
+            throw new IllegalArgumentException("vector dimensions incompatible: " + q.length + "!= " + d.length);
+        }
+        return IMPL.ipFloatByte(q, d);
+    }
+
     /**
      * AND bit count computed over signed bytes.
      * Copied from Lucene's XOR implementation

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

@@ -39,6 +39,11 @@ public float ipFloatBit(float[] q, byte[] d) {
         return ipFloatBitImpl(q, d);
     }
 
+    @Override
+    public float ipFloatByte(float[] q, byte[] d) {
+        return ipFloatByteImpl(q, d);
+    }
+
     public static int ipByteBitImpl(byte[] q, byte[] d) {
         assert q.length == d.length * Byte.SIZE;
         int acc0 = 0;
@@ -101,4 +106,12 @@ public static long ipByteBinByteImpl(byte[] q, byte[] d) {
         }
         return ret;
     }
+
+    public static float ipFloatByteImpl(float[] q, byte[] d) {
+        float ret = 0;
+        for (int i = 0; i < q.length; i++) {
+            ret += q[i] * d[i];
+        }
+        return ret;
+    }
 }

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

@@ -18,4 +18,6 @@ public interface ESVectorUtilSupport {
     int ipByteBit(byte[] q, byte[] d);
 
     float ipFloatBit(float[] q, byte[] d);
+
+    float ipFloatByte(float[] q, byte[] d);
 }

libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/PanamaESVectorUtilSupport.java

@@ -58,6 +58,11 @@ public float ipFloatBit(float[] q, byte[] d) {
         return DefaultESVectorUtilSupport.ipFloatBitImpl(q, d);
     }
 
+    @Override
+    public float ipFloatByte(float[] q, byte[] d) {
+        return DefaultESVectorUtilSupport.ipFloatByteImpl(q, d);
+    }
+
     private static final VectorSpecies<Byte> BYTE_SPECIES_128 = ByteVector.SPECIES_128;
     private static final VectorSpecies<Byte> BYTE_SPECIES_256 = ByteVector.SPECIES_256;
 

libs/simdvec/src/test/java/org/elasticsearch/simdvec/ESVectorUtilTests.java

@@ -32,11 +32,28 @@ public void testIpByteBit() {
     public void testIpFloatBit() {
         float[] q = new float[16];
         byte[] d = new byte[] { (byte) Integer.parseInt("01100010", 2), (byte) Integer.parseInt("10100111", 2) };
-        random().nextFloat();
+        for (int i = 0; i < q.length; i++) {
+            q[i] = random().nextFloat();
+        }
         float expected = q[1] + q[2] + q[6] + q[8] + q[10] + q[13] + q[14] + q[15];
         assertEquals(expected, ESVectorUtil.ipFloatBit(q, d), 1e-6);
     }
 
+    public void testIpFloatByte() {
+        float[] q = new float[16];
+        byte[] d = new byte[16];
+        for (int i = 0; i < q.length; i++) {
+            q[i] = random().nextFloat();
+        }
+        random().nextBytes(d);
+
+        float expected = 0;
+        for (int i = 0; i < q.length; i++) {
+            expected += q[i] * d[i];
+        }
+        assertEquals(expected, ESVectorUtil.ipFloatByte(q, d), 1e-6);
+    }
+
     public void testBitAndCount() {
         testBasicBitAndImpl(ESVectorUtil::andBitCountLong);
     }

modules/ingest-geoip/src/main/java/org/elasticsearch/ingest/geoip/GeoIpDownloaderTaskExecutor.java

@@ -398,7 +398,7 @@ private void stopTask(Runnable onFailure) {
             ActionListener.runAfter(listener, () -> {
                 IndexAbstraction databasesAbstraction = clusterService.state()
                     .metadata()
-                    .getProject()
+                    .getDefaultProject()
                     .getIndicesLookup()
                     .get(DATABASES_INDEX);
                 if (databasesAbstraction != null) {

modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/145_dense_vector_byte_basic.yml

@@ -107,6 +107,38 @@ setup:
   - match: {hits.hits.2._id: "1"}
   - match: {hits.hits.2._score: 1632.0}
 ---
+"Dot Product float":
+  - requires:
+      capabilities:
+        - path: /_search
+          capabilities: [byte_float_dot_product_capability]
+      test_runner_features: [capabilities]
+      reason: "float vector queries capability added"
+  - do:
+      headers:
+        Content-Type: application/json
+      search:
+        rest_total_hits_as_int: true
+        body:
+          query:
+            script_score:
+              query: {match_all: {} }
+              script:
+                source: "dotProduct(params.query_vector, 'vector')"
+                params:
+                  query_vector: [0.5, 111.3, -13.0, 14.8, -156.0]
+
+  - match: {hits.total: 3}
+
+  - match: {hits.hits.0._id: "2"}
+  - match: {hits.hits.0._score: 32865.2}
+
+  - match: {hits.hits.1._id: "3"}
+  - match: {hits.hits.1._score: 21413.4}
+
+  - match: {hits.hits.2._id: "1"}
+  - match: {hits.hits.2._score: 1862.3}
+---
 "Cosine Similarity":
   - do:
       headers:
@@ -198,3 +230,39 @@ setup:
   - match: {hits.hits.2._id: "1"}
   - gte: {hits.hits.2._score: 0.509}
   - lte: {hits.hits.2._score: 0.512}
+
+---
+"Cosine Similarity float":
+  - requires:
+      capabilities:
+        - path: /_search
+          capabilities: [byte_float_dot_product_capability]
+      test_runner_features: [capabilities]
+      reason: "float vector queries capability added"
+  - do:
+      headers:
+        Content-Type: application/json
+      search:
+        rest_total_hits_as_int: true
+        body:
+          query:
+            script_score:
+              query: {match_all: {} }
+              script:
+                source: "cosineSimilarity(params.query_vector, 'vector')"
+                params:
+                  query_vector: [0.5, 111.3, -13.0, 14.8, -156.0]
+
+  - match: {hits.total: 3}
+
+  - match: {hits.hits.0._id: "2"}
+  - gte: {hits.hits.0._score: 0.989}
+  - lte: {hits.hits.0._score: 0.992}
+
+  - match: {hits.hits.1._id: "3"}
+  - gte: {hits.hits.1._score: 0.885}
+  - lte: {hits.hits.1._score: 0.888}
+
+  - match: {hits.hits.2._id: "1"}
+  - gte: {hits.hits.2._score: 0.505}
+  - lte: {hits.hits.2._score: 0.508}

modules/parent-join/src/main/java/org/elasticsearch/join/query/ParentChildInnerHitContextBuilder.java

@@ -137,12 +137,12 @@ public TopDocsAndMaxScore topDocs(SearchHit hit) throws IOException {
                 TopDocsCollector<?> topDocsCollector;
                 MaxScoreCollector maxScoreCollector = null;
                 if (sort() != null) {
-                    topDocsCollector = new TopFieldCollectorManager(sort().sort, topN, null, Integer.MAX_VALUE, false).newCollector();
+                    topDocsCollector = new TopFieldCollectorManager(sort().sort, topN, null, Integer.MAX_VALUE).newCollector();
                     if (trackScores()) {
                         maxScoreCollector = new MaxScoreCollector();
                     }
                 } else {
-                    topDocsCollector = new TopScoreDocCollectorManager(topN, null, Integer.MAX_VALUE, false).newCollector();
+                    topDocsCollector = new TopScoreDocCollectorManager(topN, null, Integer.MAX_VALUE).newCollector();
                     maxScoreCollector = new MaxScoreCollector();
                 }
                 for (LeafReaderContext ctx : this.context.searcher().getIndexReader().leaves()) {

server/src/main/java/org/elasticsearch/index/engine/SearchBasedChangesSnapshot.java

@@ -186,13 +186,7 @@ protected TopDocs nextTopDocs() throws IOException {
         Query rangeQuery = rangeQuery(Math.max(fromSeqNo, lastSeenSeqNo), toSeqNo, indexVersionCreated);
         SortField sortBySeqNo = new SortField(SeqNoFieldMapper.NAME, SortField.Type.LONG);
 
-        TopFieldCollectorManager collectorManager = new TopFieldCollectorManager(
-            new Sort(sortBySeqNo),
-            searchBatchSize,
-            afterDoc,
-            0,
-            false
-        );
+        TopFieldCollectorManager collectorManager = new TopFieldCollectorManager(new Sort(sortBySeqNo), searchBatchSize, afterDoc, 0);
         TopDocs results = indexSearcher.search(rangeQuery, collectorManager);
 
         if (results.scoreDocs.length > 0) {