Small name refactoring, fix adjusting parameters

Author: carlosdelest

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

@@ -70,7 +70,7 @@
 import org.elasticsearch.search.vectors.ESDiversifyingChildrenFloatKnnVectorQuery;
 import org.elasticsearch.search.vectors.ESKnnByteVectorQuery;
 import org.elasticsearch.search.vectors.ESKnnFloatVectorQuery;
-import org.elasticsearch.search.vectors.KnnRescoreVectorQuery;
+import org.elasticsearch.search.vectors.RescoreKnnVectorQuery;
 import org.elasticsearch.search.vectors.VectorData;
 import org.elasticsearch.search.vectors.VectorSimilarityQuery;
 import org.elasticsearch.xcontent.ToXContent;
@@ -2096,10 +2096,12 @@ private Query createKnnByteQuery(
                 float squaredMagnitude = VectorUtil.dotProduct(queryVector, queryVector);
                 elementType.checkVectorMagnitude(similarity, ElementType.errorByteElementsAppender(queryVector), squaredMagnitude);
             }
-            Integer adjustedK = k == null || needsRescore(rescoreOversample) == false
-                ? null
-                : Math.min(OVERSAMPLE_LIMIT, (int) Math.ceil(k * rescoreOversample));
-            int adjustedNumCands = Math.max(adjustedK == null ? 0 : adjustedK, numCands);
+            Integer adjustedK = k;
+            int adjustedNumCands = numCands;
+            if (needsRescore(rescoreOversample) && adjustedK != null) {
+                adjustedK = Math.min(OVERSAMPLE_LIMIT, (int) Math.ceil(k * rescoreOversample));
+                adjustedNumCands = Math.max(adjustedK, numCands);
+            }
 
             Query knnQuery = parentFilter != null
                 ? new ESDiversifyingChildrenByteKnnVectorQuery(name(), queryVector, filter, adjustedK, adjustedNumCands, parentFilter)
@@ -2112,7 +2114,7 @@ private Query createKnnByteQuery(
                 );
             }
             if (needsRescore(rescoreOversample)) {
-                knnQuery = new KnnRescoreVectorQuery(
+                knnQuery = new RescoreKnnVectorQuery(
                     name(),
                     queryVector,
                     similarity.vectorSimilarityFunction(indexVersionCreated, ElementType.BYTE),
@@ -2148,10 +2150,12 @@ && isNotUnitVector(squaredMagnitude)) {
                 }
             }
 
-            Integer adjustedK = k == null || needsRescore(rescoreOversample) == false
-                ? k
-                : Integer.valueOf(Math.min(OVERSAMPLE_LIMIT, (int) Math.ceil(k * rescoreOversample)));
-            int adjustedNumCands = adjustedK == null ? numCands : Math.max(adjustedK, numCands);
+            Integer adjustedK = k;
+            int adjustedNumCands = numCands;
+            if (needsRescore(rescoreOversample) && adjustedK != null) {
+                adjustedK = Math.min(OVERSAMPLE_LIMIT, (int) Math.ceil(k * rescoreOversample));
+                adjustedNumCands = Math.max(adjustedK, numCands);
+            }
             Query knnQuery = parentFilter != null
                 ? new ESDiversifyingChildrenFloatKnnVectorQuery(name(), queryVector, filter, adjustedK, adjustedNumCands, parentFilter)
                 : new ESKnnFloatVectorQuery(name(), queryVector, adjustedK, adjustedNumCands, filter);
@@ -2163,7 +2167,7 @@ && isNotUnitVector(squaredMagnitude)) {
                 );
             }
             if (needsRescore(rescoreOversample)) {
-                knnQuery = new KnnRescoreVectorQuery(
+                knnQuery = new RescoreKnnVectorQuery(
                     name(),
                     queryVector,
                     similarity.vectorSimilarityFunction(indexVersionCreated, ElementType.FLOAT),

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

@@ -35,4 +35,8 @@ protected TopDocs mergeLeafResults(TopDocs[] perLeafResults) {
     public void profile(QueryProfiler queryProfiler) {
         queryProfiler.setVectorOpsCount(vectorOpsCount);
     }
+
+    public Integer kParam() {
+        return kParam;
+    }
 }

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

@@ -35,4 +35,8 @@ protected TopDocs mergeLeafResults(TopDocs[] perLeafResults) {
     public void profile(QueryProfiler queryProfiler) {
         queryProfiler.setVectorOpsCount(vectorOpsCount);
     }
+
+    public Integer kParam() {
+        return kParam;
+    }
 }

server/src/main/java/org/elasticsearch/search/vectors/KnnRescoreVectorQuery.java renamed to server/src/main/java/org/elasticsearch/search/vectors/RescoreKnnVectorQuery.java

@@ -28,44 +28,44 @@
 /**
  * Wraps a kNN vector query to rescore the results using the non-quantized vectors
  */
-public class KnnRescoreVectorQuery extends Query implements ProfilingQuery {
+public class RescoreKnnVectorQuery extends Query implements ProfilingQuery {
     private final String fieldName;
     private final byte[] byteTarget;
     private final float[] floatTarget;
     private final VectorSimilarityFunction vectorSimilarityFunction;
     private final Integer k;
-    private final Query vectorQuery;
+    private final Query innerQuery;
 
     private long vectorOpsCount;
 
-    public KnnRescoreVectorQuery(
+    public RescoreKnnVectorQuery(
         String fieldName,
         byte[] byteTarget,
         VectorSimilarityFunction vectorSimilarityFunction,
         Integer k,
-        Query vectorQuery
+        Query innerQuery
     ) {
         this.fieldName = fieldName;
         this.byteTarget = byteTarget;
         this.floatTarget = null;
         this.vectorSimilarityFunction = vectorSimilarityFunction;
         this.k = k;
-        this.vectorQuery = vectorQuery;
+        this.innerQuery = innerQuery;
     }
 
-    public KnnRescoreVectorQuery(
+    public RescoreKnnVectorQuery(
         String fieldName,
         float[] floatTarget,
         VectorSimilarityFunction vectorSimilarityFunction,
         Integer k,
-        Query vectorQuery
+        Query innerQuery
     ) {
         this.fieldName = fieldName;
         this.byteTarget = null;
         this.floatTarget = floatTarget;
         this.vectorSimilarityFunction = vectorSimilarityFunction;
         this.k = k;
-        this.vectorQuery = vectorQuery;
+        this.innerQuery = innerQuery;
     }
 
     @Override
@@ -81,7 +81,7 @@ public Query rewrite(IndexSearcher searcher) throws IOException {
         } else {
             valueSource = new VectorSimilarityFloatValueSource(fieldName, floatTarget, vectorSimilarityFunction);
         }
-        FunctionScoreQuery functionScoreQuery = new FunctionScoreQuery(vectorQuery, valueSource);
+        FunctionScoreQuery functionScoreQuery = new FunctionScoreQuery(innerQuery, valueSource);
         Query query = searcher.rewrite(functionScoreQuery);
 
         if (k == null) {
@@ -103,6 +103,14 @@ public Query rewrite(IndexSearcher searcher) throws IOException {
         return new KnnScoreDocQuery(docIds, scores, searcher.getIndexReader());
     }
 
+    public Query innerQuery() {
+        return innerQuery;
+    }
+
+    public Integer k() {
+        return k;
+    }
+
     @Override
     public void profile(QueryProfiler queryProfiler) {
         queryProfiler.setVectorOpsCount(vectorOpsCount);
@@ -119,18 +127,18 @@ public void visit(QueryVisitor visitor) {
     public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
-        KnnRescoreVectorQuery that = (KnnRescoreVectorQuery) o;
+        RescoreKnnVectorQuery that = (RescoreKnnVectorQuery) o;
         return Objects.equals(fieldName, that.fieldName)
             && Objects.deepEquals(byteTarget, that.byteTarget)
             && Objects.deepEquals(floatTarget, that.floatTarget)
             && vectorSimilarityFunction == that.vectorSimilarityFunction
             && Objects.equals(k, that.k)
-            && Objects.equals(vectorQuery, that.vectorQuery);
+            && Objects.equals(innerQuery, that.innerQuery);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(fieldName, Arrays.hashCode(byteTarget), Arrays.hashCode(floatTarget), vectorSimilarityFunction, k, vectorQuery);
+        return Objects.hash(fieldName, Arrays.hashCode(byteTarget), Arrays.hashCode(floatTarget), vectorSimilarityFunction, k, innerQuery);
     }
 
     @Override
@@ -144,7 +152,7 @@ public String toString(String field) {
         }
         sb.append(", vectorSimilarityFunction=").append(vectorSimilarityFunction);
         sb.append(", k=").append(k);
-        sb.append(", vectorQuery=").append(vectorQuery);
+        sb.append(", vectorQuery=").append(innerQuery);
         sb.append('}');
         return sb.toString();
     }