Simplify logic for RescoreKnnVectorQuery now that k is not modifiable

Author: carlosdelest

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

@@ -2134,7 +2134,6 @@ && isNotUnitVector(squaredMagnitude)) {
                     name(),
                     queryVector,
                     similarity.vectorSimilarityFunction(indexVersionCreated, ElementType.FLOAT),
-                    k,
                     knnQuery
                 );
             }

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

@@ -16,8 +16,6 @@
 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;
 
@@ -32,7 +30,6 @@ 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;
@@ -41,13 +38,11 @@ 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;
     }
 
@@ -58,34 +53,13 @@ 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);
-        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());
+        return searcher.rewrite(functionScoreQuery);
     }
 
     public Query innerQuery() {
         return innerQuery;
     }
 
-    public Integer k() {
-        return k;
-    }
-
     @Override
     public void profile(QueryProfiler queryProfiler) {
         if (innerQuery instanceof QueryProfilerProvider queryProfilerProvider) {
@@ -111,24 +85,27 @@ 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, k, innerQuery);
+        return Objects.hash(fieldName, Arrays.hashCode(floatTarget), vectorSimilarityFunction, innerQuery);
     }
 
     @Override
     public String toString(String field) {
-        final StringBuilder sb = new StringBuilder("KnnRescoreVectorQuery{");
-        sb.append("fieldName='").append(fieldName).append('\'');
-        sb.append(", floatTarget=").append(floatTarget[0]).append("...");
-        sb.append(", vectorSimilarityFunction=").append(vectorSimilarityFunction);
-        sb.append(", k=").append(k);
-        sb.append(", vectorQuery=").append(innerQuery);
-        sb.append('}');
-        return sb.toString();
+        return "KnnRescoreVectorQuery{"
+            + "fieldName='"
+            + fieldName
+            + '\''
+            + ", floatTarget="
+            + floatTarget[0]
+            + "..."
+            + ", vectorSimilarityFunction="
+            + vectorSimilarityFunction
+            + ", vectorQuery="
+            + innerQuery
+            + '}';
     }
 }

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

@@ -443,7 +443,7 @@ public void testRescoreOversampleUsedWithoutQuantization() {
         }
     }
 
-    public void testRescoreOversampleModifiesKnnParams() {
+    public void testRescoreOversampleModifiesNumCandidates() {
         DenseVectorFieldType fieldType = new DenseVectorFieldType(
             "f",
             IndexVersion.current(),
@@ -456,36 +456,38 @@ public void testRescoreOversampleModifiesKnnParams() {
         );
 
         // Total results is k, internal k is multiplied by oversample
-        checkRescoreQueryParameters(fieldType, 10, 200, 2.5F, 10, 25, 200);
+        checkRescoreQueryParameters(fieldType, 10, 200, 2.5F, 10, 500);
         // If numCands < k, update numCands to k
-        checkRescoreQueryParameters(fieldType, 10, 20, 2.5F, 10, 25, 25);
-        // Oversampling limit
-        checkRescoreQueryParameters(fieldType, 1000, 1000, 11.0F, 1000, 10000, 10000);
-        checkRescoreQueryParameters(fieldType, 5000, 7500, 2.5F, 5000, 10000, 10000);
+        checkRescoreQueryParameters(fieldType, 10, 20, 2.5F, 10, 50);
+        // Oversampling limits for num candidates
+        checkRescoreQueryParameters(fieldType, 1000, 1000, 11.0F, 1000, 10000);
+        checkRescoreQueryParameters(fieldType, 5000, 7500, 2.5F, 5000, 10000);
+        // Oversampling is capped at k as a minimum
+        checkRescoreQueryParameters(fieldType, 10, 100, 0.01F, 10, 10);
+        // Oversampling is capped at 1 as a minimum if k is not specified
+        checkRescoreQueryParameters(fieldType, null, 100, 0.0001F, null, 1);
     }
 
     private static void checkRescoreQueryParameters(
         DenseVectorFieldType fieldType,
-        int k,
+        Integer k,
         int candidates,
-        float oversample,
-        int expectedResults,
-        int expectedK,
+        float numCandsFactor,
+        Integer expectedK,
         int expectedCandidates
     ) {
         Query query = fieldType.createKnnQuery(
             VectorData.fromFloats(new float[] { 1, 4, 10 }),
             k,
             candidates,
-            oversample,
+            numCandsFactor,
             null,
             null,
             null
         );
         RescoreKnnVectorQuery rescoreQuery = (RescoreKnnVectorQuery) query;
         ESKnnFloatVectorQuery esKnnQuery = (ESKnnFloatVectorQuery) rescoreQuery.innerQuery();
-        assertThat("Unexpected total results", rescoreQuery.k(), equalTo(expectedResults));
-        assertThat("Unexpected k parameter", esKnnQuery.kParam(), equalTo(expectedK));
+        assertThat("Unexpected total results", esKnnQuery.kParam(), equalTo(expectedK));
         assertThat("Unexpected candidates", esKnnQuery.getK(), equalTo(expectedCandidates));
     }
 }

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

@@ -257,9 +257,9 @@ public void testInvalidK() {
     public void testInvalidRescoreVectorBuilder() {
         IllegalArgumentException e = expectThrows(
             IllegalArgumentException.class,
-            () -> new KnnSearchBuilder("field", randomVector(3), 0, 100, new RescoreVectorBuilder(1.0F), null)
+            () -> new KnnSearchBuilder("field", randomVector(3), 10, 100, new RescoreVectorBuilder(0.0F), null)
         );
-        assertThat(e.getMessage(), containsString("[oversample] must be > 1.0"));
+        assertThat(e.getMessage(), containsString("[num_candidates_factor] must be > 0.0"));
     }
 
     public void testRewrite() throws Exception {

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

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