Wrap ES KNN queries with PatienceKNN query

docs/changelog/127223.yaml

@@ -0,0 +1,5 @@
+pr: 127223
+summary: Wrap ES KNN queries with PatienceKNN query
+area: Vector Search
+type: feature
+issues: []

docs/reference/elasticsearch/mapping-reference/dense-vector.md

@@ -281,6 +281,9 @@ $$$dense-vector-index-options$$$
     :   In case a knn query specifies a `rescore_vector` parameter, the query `rescore_vector` parameter will be used instead.
     :   See [oversampling and rescoring quantized vectors](docs-content://solutions/search/vector/knn.md#dense-vector-knn-search-rescoring) for details.
 :::::
+
+`early_termination`
+:   (Optional, boolean) Apply _patience_ based early termination strategy to knn queries over HNSW graphs (see [paper](https://cs.uwaterloo.ca/~jimmylin/publications/Teofili_Lin_ECIR2025.pdf)). This is expected to produce Only applicable to `hnsw`, `int8_hnsw`, `int4_hnsw` and `bbq_hnsw` index types.
 ::::
 
 

qa/vector/src/main/java/org/elasticsearch/test/knn/CmdLineArgs.java

@@ -47,7 +47,8 @@ record CmdLineArgs(
     VectorSimilarityFunction vectorSpace,
     int quantizeBits,
     VectorEncoding vectorEncoding,
-    int dimensions
+    int dimensions,
+    boolean earlyTermination
 ) implements ToXContentObject {
 
     static final ParseField DOC_VECTORS_FIELD = new ParseField("doc_vectors");
@@ -70,6 +71,7 @@ record CmdLineArgs(
     static final ParseField QUANTIZE_BITS_FIELD = new ParseField("quantize_bits");
     static final ParseField VECTOR_ENCODING_FIELD = new ParseField("vector_encoding");
     static final ParseField DIMENSIONS_FIELD = new ParseField("dimensions");
+    static final ParseField EARLY_TERMINATION_FIELD = new ParseField("early_termination");
 
     static CmdLineArgs fromXContent(XContentParser parser) throws IOException {
         Builder builder = PARSER.apply(parser, null);
@@ -99,6 +101,7 @@ static CmdLineArgs fromXContent(XContentParser parser) throws IOException {
         PARSER.declareInt(Builder::setQuantizeBits, QUANTIZE_BITS_FIELD);
         PARSER.declareString(Builder::setVectorEncoding, VECTOR_ENCODING_FIELD);
         PARSER.declareInt(Builder::setDimensions, DIMENSIONS_FIELD);
+        PARSER.declareBoolean(Builder::setEarlyTermination, EARLY_TERMINATION_FIELD);
     }
 
     @Override
@@ -157,6 +160,7 @@ static class Builder {
         private int quantizeBits = 8;
         private VectorEncoding vectorEncoding = VectorEncoding.FLOAT32;
         private int dimensions;
+        private boolean earlyTermination;
 
         public Builder setDocVectors(String docVectors) {
             this.docVectors = PathUtils.get(docVectors);
@@ -258,6 +262,11 @@ public Builder setDimensions(int dimensions) {
             return this;
         }
 
+        public Builder setEarlyTermination(Boolean patience) {
+            this.earlyTermination = patience;
+            return this;
+        }
+
         public CmdLineArgs build() {
             if (docVectors == null) {
                 throw new IllegalArgumentException("Document vectors path must be provided");
@@ -287,7 +296,8 @@ public CmdLineArgs build() {
                 vectorSpace,
                 quantizeBits,
                 vectorEncoding,
-                dimensions
+                dimensions,
+                earlyTermination
             );
         }
     }

qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexTester.java

@@ -202,7 +202,7 @@ public static void main(String[] args) throws Exception {
             }
             if (cmdLineArgs.queryVectors() != null && cmdLineArgs.numQueries() > 0) {
                 KnnSearcher knnSearcher = new KnnSearcher(indexPath, cmdLineArgs);
-                knnSearcher.runSearch(result);
+                knnSearcher.runSearch(result, cmdLineArgs.earlyTermination());
             }
             formattedResults.results.add(result);
         }

qa/vector/src/main/java/org/elasticsearch/test/knn/KnnSearcher.java

@@ -33,6 +33,9 @@
 import org.apache.lucene.queries.function.valuesource.FloatKnnVectorFieldSource;
 import org.apache.lucene.queries.function.valuesource.FloatVectorSimilarityFunction;
 import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.KnnByteVectorQuery;
+import org.apache.lucene.search.KnnFloatVectorQuery;
+import org.apache.lucene.search.PatienceKnnVectorQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.search.TopDocs;
@@ -114,7 +117,7 @@ class KnnSearcher {
         this.searchThreads = cmdLineArgs.searchThreads();
     }
 
-    void runSearch(KnnIndexTester.Results finalResults) throws IOException {
+    void runSearch(KnnIndexTester.Results finalResults, boolean earlyTermination) throws IOException {
         TopDocs[] results = new TopDocs[numQueryVectors];
         int[][] resultIds = new int[numQueryVectors][];
         long elapsed, totalCpuTimeMS, totalVisited = 0;
@@ -153,10 +156,10 @@ void runSearch(KnnIndexTester.Results finalResults) throws IOException {
                     for (int i = 0; i < numQueryVectors; i++) {
                         if (vectorEncoding.equals(VectorEncoding.BYTE)) {
                             targetReader.next(targetBytes);
-                            doVectorQuery(targetBytes, searcher);
+                            doVectorQuery(targetBytes, searcher, earlyTermination);
                         } else {
                             targetReader.next(target);
-                            doVectorQuery(target, searcher);
+                            doVectorQuery(target, searcher, earlyTermination);
                         }
                     }
                     targetReader.reset();
@@ -165,10 +168,10 @@ void runSearch(KnnIndexTester.Results finalResults) throws IOException {
                     for (int i = 0; i < numQueryVectors; i++) {
                         if (vectorEncoding.equals(VectorEncoding.BYTE)) {
                             targetReader.next(targetBytes);
-                            results[i] = doVectorQuery(targetBytes, searcher);
+                            results[i] = doVectorQuery(targetBytes, searcher, earlyTermination);
                         } else {
                             targetReader.next(target);
-                            results[i] = doVectorQuery(target, searcher);
+                            results[i] = doVectorQuery(target, searcher, earlyTermination);
                         }
                     }
                     KnnIndexTester.ThreadDetails endThreadDetails = new KnnIndexTester.ThreadDetails();
@@ -263,7 +266,7 @@ private boolean isNewer(Path path, Path... others) throws IOException {
         return true;
     }
 
-    TopDocs doVectorQuery(byte[] vector, IndexSearcher searcher) throws IOException {
+    TopDocs doVectorQuery(byte[] vector, IndexSearcher searcher, boolean earlyTermination) throws IOException {
         Query knnQuery;
         if (overSamplingFactor > 1f) {
             throw new IllegalArgumentException("oversampling factor > 1 is not supported for byte vectors");
@@ -279,6 +282,9 @@ TopDocs doVectorQuery(byte[] vector, IndexSearcher searcher) throws IOException
                 null,
                 DenseVectorFieldMapper.FilterHeuristic.ACORN.getKnnSearchStrategy()
             );
+            if (indexType == KnnIndexTester.IndexType.HNSW && earlyTermination) {
+                knnQuery = PatienceKnnVectorQuery.fromByteQuery((KnnByteVectorQuery) knnQuery);
+            }
         }
         QueryProfiler profiler = new QueryProfiler();
         TopDocs docs = searcher.search(knnQuery, this.topK);
@@ -287,7 +293,7 @@ TopDocs doVectorQuery(byte[] vector, IndexSearcher searcher) throws IOException
         return new TopDocs(new TotalHits(profiler.getVectorOpsCount(), docs.totalHits.relation()), docs.scoreDocs);
     }
 
-    TopDocs doVectorQuery(float[] vector, IndexSearcher searcher) throws IOException {
+    TopDocs doVectorQuery(float[] vector, IndexSearcher searcher, boolean earlyTermination) throws IOException {
         Query knnQuery;
         int topK = this.topK;
         if (overSamplingFactor > 1f) {
@@ -306,16 +312,22 @@ TopDocs doVectorQuery(float[] vector, IndexSearcher searcher) throws IOException
                 null,
                 DenseVectorFieldMapper.FilterHeuristic.ACORN.getKnnSearchStrategy()
             );
+            if (indexType == KnnIndexTester.IndexType.HNSW && earlyTermination) {
+                knnQuery = PatienceKnnVectorQuery.fromFloatQuery((KnnFloatVectorQuery) knnQuery);
+            }
         }
         if (overSamplingFactor > 1f) {
             // oversample the topK results to get more candidates for the final result
             knnQuery = new RescoreKnnVectorQuery(VECTOR_FIELD, vector, similarityFunction, this.topK, knnQuery);
         }
         QueryProfiler profiler = new QueryProfiler();
         TopDocs docs = searcher.search(knnQuery, this.topK);
-        QueryProfilerProvider queryProfilerProvider = (QueryProfilerProvider) knnQuery;
-        queryProfilerProvider.profile(profiler);
-        return new TopDocs(new TotalHits(profiler.getVectorOpsCount(), docs.totalHits.relation()), docs.scoreDocs);
+        if (knnQuery instanceof QueryProfilerProvider queryProfilerProvider) {
+            queryProfilerProvider.profile(profiler);
+            return new TopDocs(new TotalHits(profiler.getVectorOpsCount(), docs.totalHits.relation()), docs.scoreDocs);
+        } else {
+            return docs;
+        }
     }
 
     private static float checkResults(int[][] results, int[][] nn, int topK) {

rest-api-spec/build.gradle

@@ -90,4 +90,6 @@ tasks.named("yamlRestCompatTestTransform").configure ({ task ->
   task.skipTest("indices.create/21_synthetic_source_stored/field param - keep root array", "Synthetic source keep arrays now stores leaf arrays natively")
   task.skipTest("cluster.info/30_info_thread_pool/Cluster HTTP Info", "The search_throttled thread pool has been removed")
   task.skipTest("synonyms/80_synonyms_from_index/Fail loading synonyms from index if synonyms_set doesn't exist", "Synonyms do no longer fail if the synonyms_set doesn't exist")
+  task.skipTest("search.vectors/80_dense_vector_indexed_by_default/Default index options for dense_vector", "Introduced early_termination option")
+  task.skipTest("search.vectors/80_dense_vector_indexed_by_default/Indexed by default with specified similarity and index options", "Introduced early_termination option")
 })

rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/240_knn_search_early_termination.yml

@@ -0,0 +1,168 @@
+setup:
+  - requires:
+      cluster_features: [ "search.vectors.mappers.expose_hnsw_early_termination" ]
+      reason: "requires early termination for hnsw to be exposed"
+  - do:
+      indices.create:
+        index: bbq_hnsw
+        body:
+          settings:
+            index:
+              number_of_shards: 1
+          mappings:
+            properties:
+              vector:
+                type: dense_vector
+                dims: 64
+                index: true
+                similarity: max_inner_product
+                index_options:
+                  type: bbq_hnsw
+                  early_termination: true
+
+  - do:
+      index:
+        index: bbq_hnsw
+        id: "1"
+        body:
+          vector: [0.077,  0.32 , -0.205,  0.63 ,  0.032,  0.201,  0.167, -0.313,
+                   0.176,  0.531, -0.375,  0.334, -0.046,  0.078, -0.349,  0.272,
+                   0.307, -0.083,  0.504,  0.255, -0.404,  0.289, -0.226, -0.132,
+                   -0.216,  0.49 ,  0.039,  0.507, -0.307,  0.107,  0.09 , -0.265,
+                   -0.285,  0.336, -0.272,  0.369, -0.282,  0.086, -0.132,  0.475,
+                   -0.224,  0.203,  0.439,  0.064,  0.246, -0.396,  0.297,  0.242,
+                   -0.028,  0.321, -0.022, -0.009, -0.001  ,  0.031, -0.533,  0.45,
+                   -0.683,  1.331,  0.194, -0.157, -0.1  , -0.279, -0.098, -0.176]
+  - do:
+      index:
+        index: bbq_hnsw
+        id: "2"
+        body:
+          vector: [0.196,  0.514,  0.039,  0.555, -0.042,  0.242,  0.463, -0.348,
+                   -0.08 ,  0.442, -0.067, -0.05 , -0.001,  0.298, -0.377,  0.048,
+                   0.307,  0.159,  0.278,  0.119, -0.057,  0.333, -0.289, -0.438,
+                   -0.014,  0.361, -0.169,  0.292, -0.229,  0.123,  0.031, -0.138,
+                   -0.139,  0.315, -0.216,  0.322, -0.445, -0.059,  0.071,  0.429,
+                   -0.602, -0.142,  0.11 ,  0.192,  0.259, -0.241,  0.181, -0.166,
+                   0.082,  0.107, -0.05 ,  0.155,  0.011,  0.161, -0.486,  0.569,
+                   -0.489,  0.901,  0.208,  0.011, -0.209, -0.153, -0.27 , -0.013]
+
+  - do:
+      index:
+        index: bbq_hnsw
+        id: "3"
+        body:
+          name: rabbit.jpg
+          vector: [0.139,  0.178, -0.117,  0.399,  0.014, -0.139,  0.347, -0.33 ,
+                   0.139,  0.34 , -0.052, -0.052, -0.249,  0.327, -0.288,  0.049,
+                   0.464,  0.338,  0.516,  0.247, -0.104,  0.259, -0.209, -0.246,
+                   -0.11 ,  0.323,  0.091,  0.442, -0.254,  0.195, -0.109, -0.058,
+                   -0.279,  0.402, -0.107,  0.308, -0.273,  0.019,  0.082,  0.399,
+                   -0.658, -0.03 ,  0.276,  0.041,  0.187, -0.331,  0.165,  0.017,
+                   0.171, -0.203, -0.198,  0.115, -0.007,  0.337, -0.444,  0.615,
+                   -0.657,  1.285,  0.2  , -0.062,  0.038,  0.089, -0.068, -0.058]
+  - do:
+      index:
+        index: bbq_hnsw
+        id: "4"
+        body:
+          vector: [ 0.123,  0.32 , -0.205,  -0.63 ,  0.132,  0.201,  0.167, -0.313,
+                    0.176,  0.531, -0.375,  0.334, -0.046,  0.078, -0.349,  0.272,
+                    0.307, -0.083,  0.504,  0.255, -0.904,  0.289, -0.226, -0.132,
+                    -0.216,  0.49 ,  0.039,  0.507, -0.307,  0.107,  0.09 , -0.265,
+                    -0.285,  0.336, -0.272,  0.161, -0.282,  0.086, -0.132,  0.475,
+                    -0.224,  0.203,  0.439,  0.064,  0.246, -0.396,  0.297,  0.242,
+                    -0.028,  0.321, -0.032, -0.009, -0.001  ,  0.031, -0.533,  0.45,
+                    -0.683,  1.331,  0.294, -0.157, -0.1  , -0.279, -0.098, -0.176 ]
+
+  - do:
+      index:
+        index: bbq_hnsw
+        id: "5"
+        body:
+          vector: [ 0.456,  0.32 , -0.205,  0.63 ,  0.032,  0.201,  0.167, -0.313,
+                    0.176,  0.531, -0.375,  0.334, -0.103,  0.078, -0.488,  0.272,
+                    0.307, -0.083,  0.504,  0.255, -0.404,  0.289, -0.226, -0.132,
+                    -0.216,  0.49 ,  0.039,  0.507, -0.307,  0.107,  0.084 , -0.093,
+                    -0.285,  0.336, -0.272,  0.369, -0.743,  0.086, -0.132,  0.475,
+                    -0.224,  0.203, -0.439,  0.064,  0.246, -0.396,  0.297,  0.242,
+                    -0.028,  0.321, -0.022, -0.009, -0.001  ,  0.031, -0.533,  0.45,
+                    -0.683,  1.331,  0.194, -0.157, -0.1  , -0.279, -0.098, -0.789 ]
+  - do:
+      indices.flush:
+        index: bbq_hnsw
+
+  - do:
+      indices.refresh: {}
+
+---
+"Run knn search with early termination":
+
+  - do:
+      search:
+        index: bbq_hnsw
+        body:
+          query:
+            knn:
+              field: vector
+              query_vector: [0.128,  0.067, -0.08 ,  0.395, -0.11 , -0.259,  0.473, -0.393,
+                             0.292,  0.571, -0.491,  0.444, -0.288,  0.198, -0.343,  0.015,
+                             0.232,  0.088,  0.228,  0.151, -0.136,  0.236, -0.273, -0.259,
+                             -0.217,  0.359, -0.207,  0.352, -0.142,  0.192, -0.061, -0.17 ,
+                             -0.343,  0.189, -0.221,  0.32 , -0.301, -0.1  ,  0.005,  0.232,
+                             -0.344,  0.136,  0.252,  0.157, -0.13 , -0.244,  0.193, -0.034,
+                             -0.12 , -0.193, -0.102,  0.252, -0.185, -0.167, -0.575,  0.582,
+                             -0.426,  0.983,  0.212,  0.204,  0.03 , -0.276, -0.425, -0.158]
+              num_candidates: 5
+              k: 3
+
+  - match: { hits.total.value: 3 }
+  - match: { hits.hits.0._id: "5" }
+  - match: { hits.hits.1._id: "1" }
+  - match: { hits.hits.2._id: "3" }
+
+---
+"Profile knn search with early termination":
+
+  - do:
+      search:
+        index: bbq_hnsw
+        body:
+          profile: true
+          query:
+            knn:
+              field: vector
+              query_vector: [0.128,  0.067, -0.08 ,  0.395, -0.11 , -0.259,  0.473, -0.393,
+                             0.292,  0.571, -0.491,  0.444, -0.288,  0.198, -0.343,  0.015,
+                             0.232,  0.088,  0.228,  0.151, -0.136,  0.236, -0.273, -0.259,
+                             -0.217,  0.359, -0.207,  0.352, -0.142,  0.192, -0.061, -0.17 ,
+                             -0.343,  0.189, -0.221,  0.32 , -0.301, -0.1  ,  0.005,  0.232,
+                             -0.344,  0.136,  0.252,  0.157, -0.13 , -0.244,  0.193, -0.034,
+                             -0.12 , -0.193, -0.102,  0.252, -0.185, -0.167, -0.575,  0.582,
+                             -0.426,  0.983,  0.212,  0.204,  0.03 , -0.276, -0.425, -0.158]
+              k: 3
+              num_candidates: 3
+
+  - match: { profile.shards.0.dfs.knn.0.vector_operations_count: null }
+
+  # Search with similarity to check number of operations are propagated correctly
+  - do:
+      search:
+        index: bbq_hnsw
+        body:
+          profile: true
+          knn:
+            field: vector
+            query_vector: [0.128,  0.067, -0.08 ,  0.395, -0.11 , -0.259,  0.473, -0.393,
+                           0.292,  0.571, -0.491,  0.444, -0.288,  0.198, -0.343,  0.015,
+                           0.232,  0.088,  0.228,  0.151, -0.136,  0.236, -0.273, -0.259,
+                           -0.217,  0.359, -0.207,  0.352, -0.142,  0.192, -0.061, -0.17 ,
+                           -0.343,  0.189, -0.221,  0.32 , -0.301, -0.1  ,  0.005,  0.232,
+                           -0.344,  0.136,  0.252,  0.157, -0.13 , -0.244,  0.193, -0.034,
+                           -0.12 , -0.193, -0.102,  0.252, -0.185, -0.167, -0.575,  0.582,
+                           -0.426,  0.983,  0.212,  0.204,  0.03 , -0.276, -0.425, -0.158]
+            k: 3
+            num_candidates: 3
+            similarity: 100000
+
+  - match: { profile.shards.0.dfs.knn.0.vector_operations_count: 5 }