Change API to use "rescore": {"oversample": 1.0}

Author: carlosdelest

server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RetrieverTelemetryIT.java

@@ -98,7 +98,7 @@ public void testTelemetryForRetrievers() throws IOException {
         {
             performSearch(
                 new SearchSourceBuilder().retriever(
-                    new StandardRetrieverBuilder(new KnnVectorQueryBuilder("vector", new float[] { 1.0f }, 10, 15, null))
+                    new StandardRetrieverBuilder(new KnnVectorQueryBuilder("vector", new float[] { 1.0f }, 10, 15, null, null))
                 )
             );
         }

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

@@ -2097,8 +2097,10 @@ private Query createKnnByteQuery(
                 float squaredMagnitude = VectorUtil.dotProduct(queryVector, queryVector);
                 elementType.checkVectorMagnitude(similarity, ElementType.errorByteElementsAppender(queryVector), squaredMagnitude);
             }
-            int adjustedK = rescoreOversample == null ? k : Math.min(OVERSAMPLE_LIMIT, (int) Math.ceil(k * rescoreOversample));
-            int adjustedNumCands = Math.max(adjustedK, numCands);
+            Integer adjustedK = k == null || rescoreOversample == null
+                ? null
+                : Math.min(OVERSAMPLE_LIMIT, (int) Math.ceil(k * rescoreOversample));
+            int adjustedNumCands = Math.max(adjustedK == null ? 0 : adjustedK, numCands);
 
             Query knnQuery = parentFilter != null
                 ? new ESDiversifyingChildrenByteKnnVectorQuery(name(), queryVector, filter, adjustedK, adjustedNumCands, parentFilter)
@@ -2149,8 +2151,10 @@ && isNotUnitVector(squaredMagnitude)) {
                 }
             }
 
-            int adjustedK = rescoreOversample == null ? k : Math.min(OVERSAMPLE_LIMIT, (int) Math.ceil(k * rescoreOversample));
-            int adjustedNumCands = Math.max(adjustedK, numCands);
+            Integer adjustedK = k == null || rescoreOversample == null
+                ? k
+                : Integer.valueOf(Math.min(OVERSAMPLE_LIMIT, (int) Math.ceil(k * rescoreOversample)));
+            int adjustedNumCands = adjustedK == null ? numCands : Math.max(adjustedK, numCands);
             Query knnQuery = parentFilter != null
                 ? new ESDiversifyingChildrenFloatKnnVectorQuery(name(), queryVector, filter, adjustedK, adjustedNumCands, parentFilter)
                 : new ESKnnFloatVectorQuery(name(), queryVector, adjustedK, adjustedNumCands, filter);

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

@@ -407,7 +407,7 @@ public KnnVectorQueryBuilder toQueryBuilder() {
         if (queryVectorBuilder != null) {
             throw new IllegalArgumentException("missing rewrite");
         }
-        return new KnnVectorQueryBuilder(field, queryVector, null, numCands, similarity).boost(boost)
+        return new KnnVectorQueryBuilder(field, queryVector, null, numCands, null, similarity).boost(boost)
             .queryName(queryName)
             .addFilterQueries(filterQueries);
     }

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

@@ -256,7 +256,7 @@ public KnnVectorQueryBuilder toQueryBuilder() {
             if (numCands > NUM_CANDS_LIMIT) {
                 throw new IllegalArgumentException("[" + NUM_CANDS_FIELD.getPreferredName() + "] cannot exceed [" + NUM_CANDS_LIMIT + "]");
             }
-            return new KnnVectorQueryBuilder(field, queryVector, null, numCands, null);
+            return new KnnVectorQueryBuilder(field, queryVector, null, numCands, null, null);
         }
 
         @Override

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

@@ -67,9 +67,9 @@ public class KnnVectorQueryBuilder extends AbstractQueryBuilder<KnnVectorQueryBu
     public static final ParseField NUM_CANDS_FIELD = new ParseField("num_candidates");
     public static final ParseField QUERY_VECTOR_FIELD = new ParseField("query_vector");
     public static final ParseField VECTOR_SIMILARITY_FIELD = new ParseField("similarity");
-    public static final ParseField RESCORE_VECTOR_OVERSAMPLE = new ParseField("rescore_vector_oversample");
     public static final ParseField FILTER_FIELD = new ParseField("filter");
     public static final ParseField QUERY_VECTOR_BUILDER_FIELD = new ParseField("query_vector_builder");
+    public static final ParseField RESCORE_FIELD = new ParseField("rescore");
 
     public static final ConstructingObjectParser<KnnVectorQueryBuilder, Void> PARSER = new ConstructingObjectParser<>(
         "knn",
@@ -80,8 +80,8 @@ public class KnnVectorQueryBuilder extends AbstractQueryBuilder<KnnVectorQueryBu
             null,
             (Integer) args[2],
             (Integer) args[3],
-            (Float) args[4],
-            (Float) args[6]
+            (RescoreVectorBuilder) args[6],
+            (Float) args[4]
         )
     );
 
@@ -101,7 +101,12 @@ public class KnnVectorQueryBuilder extends AbstractQueryBuilder<KnnVectorQueryBu
             (p, c, n) -> p.namedObject(QueryVectorBuilder.class, n, c),
             QUERY_VECTOR_BUILDER_FIELD
         );
-        PARSER.declareFloat(optionalConstructorArg(), RESCORE_VECTOR_OVERSAMPLE);
+        PARSER.declareField(
+            optionalConstructorArg(),
+            (p, c) -> RescoreVectorBuilder.fromXContent(p),
+            RESCORE_FIELD,
+            ObjectParser.ValueType.OBJECT_OR_NULL
+        );
         PARSER.declareFieldArray(
             KnnVectorQueryBuilder::addFilterQueries,
             (p, c) -> AbstractQueryBuilder.parseTopLevelQuery(p),
@@ -123,10 +128,17 @@ public static KnnVectorQueryBuilder fromXContent(XContentParser parser) {
     private final Float vectorSimilarity;
     private final QueryVectorBuilder queryVectorBuilder;
     private final Supplier<float[]> queryVectorSupplier;
-    private final Float rescoreOversample;
+    private final RescoreVectorBuilder rescoreVectorBuilder;
 
-    public KnnVectorQueryBuilder(String fieldName, float[] queryVector, Integer k, Integer numCands, Float vectorSimilarity) {
-        this(fieldName, VectorData.fromFloats(queryVector), null, null, k, numCands, vectorSimilarity, null);
+    public KnnVectorQueryBuilder(
+        String fieldName,
+        float[] queryVector,
+        Integer k,
+        Integer numCands,
+        RescoreVectorBuilder rescoreVectorBuilder,
+        Float vectorSimilarity
+    ) {
+        this(fieldName, VectorData.fromFloats(queryVector), null, null, k, numCands, rescoreVectorBuilder, vectorSimilarity);
     }
 
     public KnnVectorQueryBuilder(
@@ -136,27 +148,29 @@ public KnnVectorQueryBuilder(
         Integer numCands,
         Float vectorSimilarity
     ) {
-        this(fieldName, null, queryVectorBuilder, null, k, numCands, vectorSimilarity, null);
+        this(fieldName, null, queryVectorBuilder, null, k, numCands, null, vectorSimilarity);
     }
 
-    public KnnVectorQueryBuilder(String fieldName, byte[] queryVector, Integer k, Integer numCands, Float vectorSimilarity) {
-        this(fieldName, VectorData.fromBytes(queryVector), null, null, k, numCands, vectorSimilarity);
-    }
-
-    public KnnVectorQueryBuilder(String fieldName, VectorData queryVector, Integer k, Integer numCands, Float vectorSimilarity) {
-        this(fieldName, queryVector, null, null, k, numCands, vectorSimilarity);
+    public KnnVectorQueryBuilder(
+        String fieldName,
+        byte[] queryVector,
+        Integer k,
+        Integer numCands,
+        RescoreVectorBuilder rescoreVectorBuilder,
+        Float vectorSimilarity
+    ) {
+        this(fieldName, VectorData.fromBytes(queryVector), null, null, k, numCands, rescoreVectorBuilder, vectorSimilarity);
     }
 
-    private KnnVectorQueryBuilder(
+    public KnnVectorQueryBuilder(
         String fieldName,
         VectorData queryVector,
-        QueryVectorBuilder queryVectorBuilder,
-        Supplier<float[]> queryVectorSupplier,
         Integer k,
         Integer numCands,
+        RescoreVectorBuilder rescoreVectorBuilder,
         Float vectorSimilarity
     ) {
-        this(fieldName, queryVector, null, null, k, numCands, vectorSimilarity, null);
+        this(fieldName, queryVector, null, null, k, numCands, rescoreVectorBuilder, vectorSimilarity);
     }
 
     private KnnVectorQueryBuilder(
@@ -166,8 +180,8 @@ private KnnVectorQueryBuilder(
         Supplier<float[]> queryVectorSupplier,
         Integer k,
         Integer numCands,
-        Float vectorSimilarity,
-        Float rescoreOversample
+        RescoreVectorBuilder rescoreVectorBuilder,
+        Float vectorSimilarity
     ) {
         if (k != null && k < 1) {
             throw new IllegalArgumentException("[" + K_FIELD.getPreferredName() + "] must be greater than 0");
@@ -204,7 +218,7 @@ private KnnVectorQueryBuilder(
         this.vectorSimilarity = vectorSimilarity;
         this.queryVectorBuilder = queryVectorBuilder;
         this.queryVectorSupplier = queryVectorSupplier;
-        this.rescoreOversample = rescoreOversample;
+        this.rescoreVectorBuilder = rescoreVectorBuilder;
     }
 
     public KnnVectorQueryBuilder(StreamInput in) throws IOException {
@@ -246,9 +260,9 @@ public KnnVectorQueryBuilder(StreamInput in) throws IOException {
             this.queryVectorBuilder = null;
         }
         if (in.getTransportVersion().onOrAfter(KNN_QUERY_RESCORE_OVERSAMPLE)) {
-            this.rescoreOversample = in.readOptionalFloat();
+            this.rescoreVectorBuilder = in.readOptional(RescoreVectorBuilder::new);
         } else {
-            this.rescoreOversample = null;
+            this.rescoreVectorBuilder = null;
         }
 
         this.queryVectorSupplier = null;
@@ -276,10 +290,6 @@ public Integer numCands() {
         return numCands;
     }
 
-    public Float rescoreOversample() {
-        return rescoreOversample;
-    }
-
     public List<QueryBuilder> filterQueries() {
         return filterQueries;
     }
@@ -289,6 +299,10 @@ public QueryVectorBuilder queryVectorBuilder() {
         return queryVectorBuilder;
     }
 
+    public RescoreVectorBuilder rescoreVectorBuilder() {
+        return rescoreVectorBuilder;
+    }
+
     public KnnVectorQueryBuilder addFilterQuery(QueryBuilder filterQuery) {
         Objects.requireNonNull(filterQuery);
         this.filterQueries.add(filterQuery);
@@ -356,7 +370,7 @@ protected void doWriteTo(StreamOutput out) throws IOException {
             out.writeOptionalNamedWriteable(queryVectorBuilder);
         }
         if (out.getTransportVersion().onOrAfter(KNN_QUERY_RESCORE_OVERSAMPLE)) {
-            out.writeOptionalFloat(rescoreOversample);
+            out.writeOptionalWriteable(rescoreVectorBuilder);
         }
     }
 
@@ -391,6 +405,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
             }
             builder.endArray();
         }
+        if (rescoreVectorBuilder != null) {
+            builder.startObject(RESCORE_FIELD.getPreferredName());
+            rescoreVectorBuilder.toXContent(builder, params);
+            builder.endObject();
+        }
         boostAndQueryNameToXContent(builder);
         builder.endObject();
     }
@@ -406,7 +425,8 @@ protected QueryBuilder doRewrite(QueryRewriteContext ctx) throws IOException {
             if (queryVectorSupplier.get() == null) {
                 return this;
             }
-            return new KnnVectorQueryBuilder(fieldName, queryVectorSupplier.get(), k, numCands, vectorSimilarity).boost(boost)
+            return new KnnVectorQueryBuilder(fieldName, queryVectorSupplier.get(), k, numCands, rescoreVectorBuilder, vectorSimilarity)
+                .boost(boost)
                 .queryName(queryName)
                 .addFilterQueries(filterQueries);
         }
@@ -428,9 +448,16 @@ protected QueryBuilder doRewrite(QueryRewriteContext ctx) throws IOException {
                 }
                 ll.onResponse(null);
             })));
-            return new KnnVectorQueryBuilder(fieldName, queryVector, queryVectorBuilder, toSet::get, k, numCands, vectorSimilarity).boost(
-                boost
-            ).queryName(queryName).addFilterQueries(filterQueries);
+            return new KnnVectorQueryBuilder(
+                fieldName,
+                queryVector,
+                queryVectorBuilder,
+                toSet::get,
+                k,
+                numCands,
+                rescoreVectorBuilder,
+                vectorSimilarity
+            ).boost(boost).queryName(queryName).addFilterQueries(filterQueries);
         }
         if (ctx.convertToInnerHitsRewriteContext() != null) {
             return new ExactKnnQueryBuilder(queryVector, fieldName, vectorSimilarity).boost(boost).queryName(queryName);
@@ -448,10 +475,16 @@ protected QueryBuilder doRewrite(QueryRewriteContext ctx) throws IOException {
             rewrittenQueries.add(rewrittenQuery);
         }
         if (changed) {
-            return new KnnVectorQueryBuilder(fieldName, queryVector, queryVectorBuilder, queryVectorSupplier, k, numCands, vectorSimilarity)
-                .boost(boost)
-                .queryName(queryName)
-                .addFilterQueries(rewrittenQueries);
+            return new KnnVectorQueryBuilder(
+                fieldName,
+                queryVector,
+                queryVectorBuilder,
+                queryVectorSupplier,
+                k,
+                numCands,
+                rescoreVectorBuilder,
+                vectorSimilarity
+            ).boost(boost).queryName(queryName).addFilterQueries(rewrittenQueries);
         }
         return this;
     }
@@ -495,6 +528,7 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
 
         DenseVectorFieldType vectorFieldType = (DenseVectorFieldType) fieldType;
         String parentPath = context.nestedLookup().getNestedParent(fieldName);
+        Float rescoreOversample = rescoreVectorBuilder() == null ? null : rescoreVectorBuilder.oversample();
 
         if (parentPath != null) {
             final BitSetProducer parentBitSet;
@@ -550,7 +584,7 @@ protected int doHashCode() {
             filterQueries,
             vectorSimilarity,
             queryVectorBuilder,
-            rescoreOversample
+            rescoreVectorBuilder
         );
     }
 
@@ -563,7 +597,7 @@ protected boolean doEquals(KnnVectorQueryBuilder other) {
             && Objects.equals(filterQueries, other.filterQueries)
             && Objects.equals(vectorSimilarity, other.vectorSimilarity)
             && Objects.equals(queryVectorBuilder, other.queryVectorBuilder)
-            && Objects.equals(rescoreOversample, other.rescoreOversample);
+            && Objects.equals(rescoreVectorBuilder, other.rescoreVectorBuilder);
     }
 
     @Override

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

@@ -0,0 +1,87 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.search.vectors;
+
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.xcontent.ConstructingObjectParser;
+import org.elasticsearch.xcontent.ParseField;
+import org.elasticsearch.xcontent.ToXContentObject;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentParser;
+
+import java.io.IOException;
+import java.util.Objects;
+
+public class RescoreVectorBuilder implements Writeable, ToXContentObject {
+
+    public static final ParseField OVERSAMPLE_FIELD = new ParseField("oversample");
+    public static final int MIN_OVERSAMPLE = 1;
+    private static final ConstructingObjectParser<RescoreVectorBuilder, Void> PARSER = new ConstructingObjectParser<>(
+        "rescore",
+        args -> new RescoreVectorBuilder((Float) args[0])
+    );
+
+    static {
+        PARSER.declareFloat(ConstructingObjectParser.optionalConstructorArg(), OVERSAMPLE_FIELD);
+    }
+
+    // Oversample is required as of now as it is the only field in the rescore vector
+    // that may change in the future, so we treat it as optional
+    private final Float oversample;
+
+    public RescoreVectorBuilder(Float oversample) {
+        Objects.requireNonNull(oversample, "[" + OVERSAMPLE_FIELD.getPreferredName() + "] must be set");
+        if (oversample <= MIN_OVERSAMPLE) {
+            throw new IllegalArgumentException("[" + OVERSAMPLE_FIELD.getPreferredName() + "] must be > " + MIN_OVERSAMPLE);
+        }
+        this.oversample = oversample;
+    }
+
+    public RescoreVectorBuilder(StreamInput in) throws IOException {
+        this.oversample = in.readOptionalFloat();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeOptionalFloat(oversample);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        if (oversample != null) {
+            builder.field(OVERSAMPLE_FIELD.getPreferredName(), oversample);
+        }
+
+        return builder;
+    }
+
+    public static RescoreVectorBuilder fromXContent(XContentParser parser) {
+        return PARSER.apply(parser, null);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        RescoreVectorBuilder that = (RescoreVectorBuilder) o;
+        return Objects.equals(oversample, that.oversample);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hashCode(oversample);
+    }
+
+    public Float oversample() {
+        return oversample;
+    }
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapper.java

@@ -556,7 +556,7 @@ public QueryBuilder semanticQuery(InferenceResults inferenceResults, float boost
                             );
                         }
 
-                        yield new KnnVectorQueryBuilder(inferenceResultsFieldName, inference, null, null, null);
+                        yield new KnnVectorQueryBuilder(inferenceResultsFieldName, inference, null, null, null,null);
                     }
                     default -> throw new IllegalStateException(
                         "Field ["