Vector similarity needs to wrap the new rescoring query and not the other way round

Author: carlosdelest

server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java

@@ -2111,29 +2111,33 @@ && isNotUnitVector(squaredMagnitude)) {
                 }
             }
 
+            Integer adjustedK = k;
             int adjustedNumCands = numCands;
             if (needsRescore(numCandsFactor)) {
-                // k <= numCands * numCandsFactor <= NUM_CANDS_OVERSAMPLE_LIMIT. Adjust otherwise.
+                // Get all candidates, get top k as part of rescoring
+                adjustedK = null;
+                // numCands * numCandsFactor <= NUM_CANDS_OVERSAMPLE_LIMIT. Adjust otherwise.
                 adjustedNumCands = Math.min((int) Math.ceil(numCands * numCandsFactor), NUM_CANDS_OVERSAMPLE_LIMIT);
             }
             Query knnQuery = parentFilter != null
-                ? new ESDiversifyingChildrenFloatKnnVectorQuery(name(), queryVector, filter, k, adjustedNumCands, parentFilter)
-                : new ESKnnFloatVectorQuery(name(), queryVector, k, adjustedNumCands, filter);
-            if (similarityThreshold != null) {
-                knnQuery = new VectorSimilarityQuery(
-                    knnQuery,
-                    similarityThreshold,
-                    similarity.score(similarityThreshold, elementType, dims)
-                );
-            }
+                ? new ESDiversifyingChildrenFloatKnnVectorQuery(name(), queryVector, filter, adjustedK, adjustedNumCands, parentFilter)
+                : new ESKnnFloatVectorQuery(name(), queryVector, adjustedK, adjustedNumCands, filter);
             if (needsRescore(numCandsFactor)) {
                 knnQuery = new RescoreKnnVectorQuery(
                     name(),
                     queryVector,
                     similarity.vectorSimilarityFunction(indexVersionCreated, ElementType.FLOAT),
+                    k,
                     knnQuery
                 );
             }
+            if (similarityThreshold != null) {
+                knnQuery = new VectorSimilarityQuery(
+                    knnQuery,
+                    similarityThreshold,
+                    similarity.score(similarityThreshold, elementType, dims)
+                );
+            }
             return knnQuery;
         }
 

server/src/main/java/org/elasticsearch/search/vectors/RescoreKnnVectorQuery.java

@@ -16,6 +16,8 @@
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryVisitor;
+import org.apache.lucene.search.ScoreDoc;
+import org.apache.lucene.search.TopDocs;
 import org.elasticsearch.index.mapper.vectors.VectorSimilarityFloatValueSource;
 import org.elasticsearch.search.profile.query.QueryProfiler;
 
@@ -30,6 +32,7 @@ public class RescoreKnnVectorQuery extends Query implements QueryProfilerProvide
     private final String fieldName;
     private final float[] floatTarget;
     private final VectorSimilarityFunction vectorSimilarityFunction;
+    private final Integer k;
     private final Query innerQuery;
 
     private QueryProfilerProvider vectorProfiling;
@@ -38,11 +41,13 @@ public RescoreKnnVectorQuery(
         String fieldName,
         float[] floatTarget,
         VectorSimilarityFunction vectorSimilarityFunction,
+        Integer k,
         Query innerQuery
     ) {
         this.fieldName = fieldName;
         this.floatTarget = floatTarget;
         this.vectorSimilarityFunction = vectorSimilarityFunction;
+        this.k = k;
         this.innerQuery = innerQuery;
     }
 
@@ -53,13 +58,34 @@ public Query rewrite(IndexSearcher searcher) throws IOException {
         // to calculate top k and return directly the query to understand how many comparisons were done
         vectorProfiling = (QueryProfilerProvider) valueSource;
         FunctionScoreQuery functionScoreQuery = new FunctionScoreQuery(innerQuery, valueSource);
-        return searcher.rewrite(functionScoreQuery);
+        Query query = searcher.rewrite(functionScoreQuery);
+
+        if (k == null) {
+            // No need to calculate top k - let the request size limit the results.
+            return query;
+        }
+
+        // Retrieve top k documents from the rescored query
+        TopDocs topDocs = searcher.search(query, k);
+        ScoreDoc[] scoreDocs = topDocs.scoreDocs;
+        int[] docIds = new int[scoreDocs.length];
+        float[] scores = new float[scoreDocs.length];
+        for (int i = 0; i < scoreDocs.length; i++) {
+            docIds[i] = scoreDocs[i].doc;
+            scores[i] = scoreDocs[i].score;
+        }
+
+        return new KnnScoreDocQuery(docIds, scores, searcher.getIndexReader());
     }
 
     public Query innerQuery() {
         return innerQuery;
     }
 
+    public Integer k() {
+        return k;
+    }
+
     @Override
     public void profile(QueryProfiler queryProfiler) {
         if (innerQuery instanceof QueryProfilerProvider queryProfilerProvider) {
@@ -85,27 +111,22 @@ public boolean equals(Object o) {
         return Objects.equals(fieldName, that.fieldName)
             && Objects.deepEquals(floatTarget, that.floatTarget)
             && vectorSimilarityFunction == that.vectorSimilarityFunction
+            && Objects.equals(k, that.k)
             && Objects.equals(innerQuery, that.innerQuery);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(fieldName, Arrays.hashCode(floatTarget), vectorSimilarityFunction, innerQuery);
+        return Objects.hash(fieldName, Arrays.hashCode(floatTarget), vectorSimilarityFunction, k, innerQuery);
     }
 
     @Override
     public String toString(String field) {
-        return "KnnRescoreVectorQuery{"
-            + "fieldName='"
-            + fieldName
-            + '\''
-            + ", floatTarget="
-            + floatTarget[0]
-            + "..."
-            + ", vectorSimilarityFunction="
-            + vectorSimilarityFunction
-            + ", vectorQuery="
-            + innerQuery
-            + '}';
+        return "KnnRescoreVectorQuery{" + "fieldName='" + fieldName + '\'' +
+            ", floatTarget=" + floatTarget[0] + "..." +
+            ", vectorSimilarityFunction=" + vectorSimilarityFunction +
+            ", k=" + k +
+            ", vectorQuery=" + innerQuery +
+            '}';
     }
 }

server/src/main/java/org/elasticsearch/search/vectors/VectorSimilarityQuery.java

@@ -29,7 +29,8 @@
 import static org.elasticsearch.common.Strings.format;
 
 /**
- * This query provides a simple post-filter for the provided Query. The query is assumed to be a Knn(Float|Byte)VectorQuery.
+ * This query provides a simple post-filter for the provided Query to limit the results of the inner query to those that have a similarity
+ * above a certain threshold
  */
 public class VectorSimilarityQuery extends Query implements QueryProfilerProvider {
     private final float similarity;

server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java

@@ -456,12 +456,12 @@ public void testRescoreOversampleModifiesNumCandidates() {
         );
 
         // Total results is k, internal k is multiplied by oversample
-        checkRescoreQueryParameters(fieldType, 10, 200, 2.5F, 10, 500);
+        checkRescoreQueryParameters(fieldType, 10, 200, 2.5F, null, 500, 10);
         // If numCands < k, update numCands to k
-        checkRescoreQueryParameters(fieldType, 10, 20, 2.5F, 10, 50);
+        checkRescoreQueryParameters(fieldType, 10, 20, 2.5F, null, 50, 10);
         // Oversampling limits for num candidates
-        checkRescoreQueryParameters(fieldType, 1000, 1000, 11.0F, 1000, 10000);
-        checkRescoreQueryParameters(fieldType, 5000, 7500, 2.5F, 5000, 10000);
+        checkRescoreQueryParameters(fieldType, 1000, 1000, 11.0F, null, 10000, 1000);
+        checkRescoreQueryParameters(fieldType, 5000, 7500, 2.5F, null, 10000, 5000);
     }
 
     private static void checkRescoreQueryParameters(
@@ -470,7 +470,8 @@ private static void checkRescoreQueryParameters(
         int candidates,
         float numCandsFactor,
         Integer expectedK,
-        int expectedCandidates
+        int expectedCandidates,
+        int expectedResults
     ) {
         Query query = fieldType.createKnnQuery(
             VectorData.fromFloats(new float[] { 1, 4, 10 }),
@@ -483,7 +484,8 @@ private static void checkRescoreQueryParameters(
         );
         RescoreKnnVectorQuery rescoreQuery = (RescoreKnnVectorQuery) query;
         ESKnnFloatVectorQuery esKnnQuery = (ESKnnFloatVectorQuery) rescoreQuery.innerQuery();
-        assertThat("Unexpected total results", esKnnQuery.kParam(), equalTo(expectedK));
+        assertThat("Unexpected total results", rescoreQuery.k(), equalTo(expectedResults));
+        assertThat("Unexpected k parameter", esKnnQuery.kParam(), equalTo(expectedK));
         assertThat("Unexpected candidates", esKnnQuery.getK(), equalTo(expectedCandidates));
     }
 }

server/src/test/java/org/elasticsearch/search/vectors/AbstractKnnVectorQueryBuilderTestCase.java

@@ -171,23 +171,18 @@ protected RescoreVectorBuilder randomRescoreVectorBuilder() {
 
     @Override
     protected void doAssertLuceneQuery(KnnVectorQueryBuilder queryBuilder, Query query, SearchExecutionContext context) throws IOException {
+        if (queryBuilder.getVectorSimilarity() != null) {
+            assertTrue(query instanceof VectorSimilarityQuery);
+            assertThat(((VectorSimilarityQuery) query).getSimilarity(), equalTo(queryBuilder.getVectorSimilarity()));
+            query = ((VectorSimilarityQuery) query).getInnerKnnQuery();
+        }
         if (queryBuilder.rescoreVectorBuilder() != null && isQuantizedElementType()) {
             RescoreKnnVectorQuery rescoreQuery = (RescoreKnnVectorQuery) query;
             query = rescoreQuery.innerQuery();
         }
-        if (queryBuilder.getVectorSimilarity() != null) {
-            assertTrue(query instanceof VectorSimilarityQuery);
-            Query knnQuery = ((VectorSimilarityQuery) query).getInnerKnnQuery();
-            assertThat(((VectorSimilarityQuery) query).getSimilarity(), equalTo(queryBuilder.getVectorSimilarity()));
-            switch (elementType()) {
-                case FLOAT -> assertTrue(knnQuery instanceof ESKnnFloatVectorQuery);
-                case BYTE -> assertTrue(knnQuery instanceof ESKnnByteVectorQuery);
-            }
-        } else {
-            switch (elementType()) {
-                case FLOAT -> assertTrue(query instanceof ESKnnFloatVectorQuery);
-                case BYTE -> assertTrue(query instanceof ESKnnByteVectorQuery);
-            }
+        switch (elementType()) {
+            case FLOAT -> assertTrue(query instanceof ESKnnFloatVectorQuery);
+            case BYTE -> assertTrue(query instanceof ESKnnByteVectorQuery);
         }
 
         BooleanQuery.Builder builder = new BooleanQuery.Builder();

server/src/test/java/org/elasticsearch/search/vectors/RescoreKnnVectorQueryTests.java

@@ -77,6 +77,7 @@ public void testRescoreDocs() throws Exception {
                     FIELD_NAME,
                     queryVector,
                     VectorSimilarityFunction.COSINE,
+                    adjustedK,
                     new MatchAllDocsQuery()
                 );
 
@@ -142,6 +143,7 @@ private void checkProfiling(float[] queryVector, IndexReader reader, Query inner
             FIELD_NAME,
             queryVector,
             VectorSimilarityFunction.COSINE,
+            k,
             innerQuery
         );
         IndexSearcher searcher = newSearcher(reader, true, false);