From 0eefe04197ae650a533ebaad01172df9ea478300 Mon Sep 17 00:00:00 2001
From: Benjamin Trent <ben.w.trent@gmail.com>
Date: Mon, 29 Jul 2024 16:01:56 -0400
Subject: [PATCH 1/4] Ensure vector similarity correctly limits inner_hits
 returned for nested kNN (#111363)

For nested kNN we support not only similarity thresholds, but also
multi-passage search while retrieving more than one nearest passage.

However, the inner_hits retrieved for the kNN search would ignore the
restricted similarity. Meaning, the inner hits would return all
passages, not just the ones within the limited similarity and this is
confusing.

closes: https://github.com/elastic/elasticsearch/issues/111093
(cherry picked from commit 69c96974de548ee3bfbfed482f0c205e18d42c8d)
---
 docs/changelog/111363.yaml                    |  6 +++
 .../search.vectors/100_knn_nested_search.yml  | 50 +++++++++++++++++++
 .../org/elasticsearch/TransportVersions.java  |  2 +
 .../action/search/DfsQueryPhase.java          |  3 +-
 .../vectors/DenseVectorFieldMapper.java       |  8 ++-
 .../search/vectors/ExactKnnQueryBuilder.java  | 37 ++++++++------
 .../vectors/KnnScoreDocQueryBuilder.java      | 37 ++++++++------
 .../search/vectors/KnnSearchBuilder.java      |  4 ++
 .../search/vectors/KnnVectorQueryBuilder.java |  2 +-
 .../action/search/DfsQueryPhaseTests.java     |  6 ++-
 .../vectors/DenseVectorFieldTypeTests.java    |  4 +-
 ...AbstractKnnVectorQueryBuilderTestCase.java | 17 +++++++
 .../vectors/ExactKnnQueryBuilderTests.java    | 23 ++++++++-
 .../vectors/KnnScoreDocQueryBuilderTests.java | 27 +++++++---
 14 files changed, 182 insertions(+), 44 deletions(-)
 create mode 100644 docs/changelog/111363.yaml

diff --git a/docs/changelog/111363.yaml b/docs/changelog/111363.yaml
new file mode 100644
index 0000000000000..2cb3c5342ea5c
--- /dev/null
+++ b/docs/changelog/111363.yaml
@@ -0,0 +1,6 @@
+pr: 111363
+summary: Ensure vector similarity correctly limits `inner_hits` returned for nested
+  kNN
+area: Vector Search
+type: bug
+issues: [111093]
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/100_knn_nested_search.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/100_knn_nested_search.yml
index 72c6abab22600..d627be2fb15c3 100644
--- a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/100_knn_nested_search.yml
+++ b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/100_knn_nested_search.yml
@@ -411,3 +411,53 @@ setup:
 
   - match: {hits.total.value: 1}
   - match: {hits.hits.0._id: "2"}
+---
+"nested Knn search with required similarity appropriately filters inner_hits":
+  - requires:
+      cluster_features: "gte_v8.16.0"
+      reason: 'bugfix for 8.16'
+
+  - do:
+      search:
+        index: test
+        body:
+          query:
+            nested:
+              path: nested
+              inner_hits:
+                size: 3
+                _source: false
+                fields:
+                  - nested.paragraph_id
+              query:
+                knn:
+                  field: nested.vector
+                  query_vector: [-0.5, 90.0, -10, 14.8, -156.0]
+                  num_candidates: 3
+                  similarity: 10.5
+
+  - match: {hits.total.value: 1}
+  - match: {hits.hits.0._id: "2"}
+  - length: {hits.hits.0.inner_hits.nested.hits.hits: 1}
+  - match: {hits.hits.0.inner_hits.nested.hits.hits.0.fields.nested.0.paragraph_id.0: "0"}
+
+  - do:
+      search:
+        index: test
+        body:
+          knn:
+            field: nested.vector
+            query_vector: [-0.5, 90.0, -10, 14.8, -156.0]
+            num_candidates: 3
+            k: 3
+            similarity: 10.5
+            inner_hits:
+              size: 3
+              _source: false
+              fields:
+                - nested.paragraph_id
+
+  - match: {hits.total.value: 1}
+  - match: {hits.hits.0._id: "2"}
+  - length: {hits.hits.0.inner_hits.nested.hits.hits: 1}
+  - match: {hits.hits.0.inner_hits.nested.hits.hits.0.fields.nested.0.paragraph_id.0: "0"}
diff --git a/server/src/main/java/org/elasticsearch/TransportVersions.java b/server/src/main/java/org/elasticsearch/TransportVersions.java
index 7ac9b1a3f8013..f2da6aba50da9 100644
--- a/server/src/main/java/org/elasticsearch/TransportVersions.java
+++ b/server/src/main/java/org/elasticsearch/TransportVersions.java
@@ -210,6 +210,8 @@ static TransportVersion def(int id) {
     public static final TransportVersion VERSIONED_MASTER_NODE_REQUESTS = def(8_701_00_0);
     public static final TransportVersion ML_INFERENCE_AMAZON_BEDROCK_ADDED = def(8_702_00_0);
     public static final TransportVersion ENTERPRISE_GEOIP_DOWNLOADER_BACKPORT_8_15 = def(8_702_00_1);
+    public static final TransportVersion FIX_VECTOR_SIMILARITY_INNER_HITS_BACKPORT_8_15 = def(8_702_00_2);
+    public static final TransportVersion FIX_VECTOR_SIMILARITY_INNER_HITS = def(8_713_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,
diff --git a/server/src/main/java/org/elasticsearch/action/search/DfsQueryPhase.java b/server/src/main/java/org/elasticsearch/action/search/DfsQueryPhase.java
index 9ddac7f13eb51..7a33eaa59eb03 100644
--- a/server/src/main/java/org/elasticsearch/action/search/DfsQueryPhase.java
+++ b/server/src/main/java/org/elasticsearch/action/search/DfsQueryPhase.java
@@ -155,7 +155,8 @@ ShardSearchRequest rewriteShardSearchRequest(ShardSearchRequest request) {
             QueryBuilder query = new KnnScoreDocQueryBuilder(
                 scoreDocs.toArray(Lucene.EMPTY_SCORE_DOCS),
                 source.knnSearch().get(i).getField(),
-                source.knnSearch().get(i).getQueryVector()
+                source.knnSearch().get(i).getQueryVector(),
+                source.knnSearch().get(i).getSimilarity()
             ).boost(source.knnSearch().get(i).boost()).queryName(source.knnSearch().get(i).queryName());
             if (nestedPath != null) {
                 query = new NestedQueryBuilder(nestedPath, query, ScoreMode.Max).innerHit(source.knnSearch().get(i).innerHit());
diff --git a/server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java
index d27c0acdb6b2e..b9912b3e097a2 100644
--- a/server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java
+++ b/server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java
@@ -1711,17 +1711,21 @@ public Query termQuery(Object value, SearchExecutionContext context) {
             throw new IllegalArgumentException("Field [" + name() + "] of type [" + typeName() + "] doesn't support term queries");
         }
 
-        public Query createExactKnnQuery(VectorData queryVector) {
+        public Query createExactKnnQuery(VectorData queryVector, Float vectorSimilarity) {
             if (isIndexed() == false) {
                 throw new IllegalArgumentException(
                     "to perform knn search on field [" + name() + "], its mapping must have [index] set to [true]"
                 );
             }
-            return switch (elementType) {
+            Query knnQuery = switch (elementType) {
                 case BYTE -> createExactKnnByteQuery(queryVector.asByteVector());
                 case FLOAT -> createExactKnnFloatQuery(queryVector.asFloatVector());
                 case BIT -> createExactKnnBitQuery(queryVector.asByteVector());
             };
+            if (vectorSimilarity != null) {
+                knnQuery = new VectorSimilarityQuery(knnQuery, vectorSimilarity, similarity.score(vectorSimilarity, elementType, dims));
+            }
+            return knnQuery;
         }
 
         private Query createExactKnnBitQuery(byte[] queryVector) {
diff --git a/server/src/main/java/org/elasticsearch/search/vectors/ExactKnnQueryBuilder.java b/server/src/main/java/org/elasticsearch/search/vectors/ExactKnnQueryBuilder.java
index 1f05b215699b1..ebefd7d8b85b6 100644
--- a/server/src/main/java/org/elasticsearch/search/vectors/ExactKnnQueryBuilder.java
+++ b/server/src/main/java/org/elasticsearch/search/vectors/ExactKnnQueryBuilder.java
@@ -32,6 +32,7 @@ public class ExactKnnQueryBuilder extends AbstractQueryBuilder<ExactKnnQueryBuil
     public static final String NAME = "exact_knn";
     private final String field;
     private final VectorData query;
+    private final Float vectorSimilarity;
 
     /**
      * Creates a query builder.
@@ -39,19 +40,10 @@ public class ExactKnnQueryBuilder extends AbstractQueryBuilder<ExactKnnQueryBuil
      * @param query    the query vector
      * @param field    the field that was used for the kNN query
      */
-    public ExactKnnQueryBuilder(float[] query, String field) {
-        this(VectorData.fromFloats(query), field);
-    }
-
-    /**
-     * Creates a query builder.
-     *
-     * @param query    the query vector
-     * @param field    the field that was used for the kNN query
-     */
-    public ExactKnnQueryBuilder(VectorData query, String field) {
+    public ExactKnnQueryBuilder(VectorData query, String field, Float vectorSimilarity) {
         this.query = query;
         this.field = field;
+        this.vectorSimilarity = vectorSimilarity;
     }
 
     public ExactKnnQueryBuilder(StreamInput in) throws IOException {
@@ -62,6 +54,12 @@ public ExactKnnQueryBuilder(StreamInput in) throws IOException {
             this.query = VectorData.fromFloats(in.readFloatArray());
         }
         this.field = in.readString();
+        if (in.getTransportVersion().onOrAfter(TransportVersions.FIX_VECTOR_SIMILARITY_INNER_HITS)
+            || in.getTransportVersion().isPatchFrom(TransportVersions.FIX_VECTOR_SIMILARITY_INNER_HITS_BACKPORT_8_15)) {
+            this.vectorSimilarity = in.readOptionalFloat();
+        } else {
+            this.vectorSimilarity = null;
+        }
     }
 
     String getField() {
@@ -72,6 +70,10 @@ VectorData getQuery() {
         return query;
     }
 
+    Float vectorSimilarity() {
+        return vectorSimilarity;
+    }
+
     @Override
     public String getWriteableName() {
         return NAME;
@@ -85,6 +87,10 @@ protected void doWriteTo(StreamOutput out) throws IOException {
             out.writeFloatArray(query.asFloatVector());
         }
         out.writeString(field);
+        if (out.getTransportVersion().onOrAfter(TransportVersions.FIX_VECTOR_SIMILARITY_INNER_HITS)
+            || out.getTransportVersion().isPatchFrom(TransportVersions.FIX_VECTOR_SIMILARITY_INNER_HITS_BACKPORT_8_15)) {
+            out.writeOptionalFloat(vectorSimilarity);
+        }
     }
 
     @Override
@@ -92,6 +98,9 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
         builder.startObject(NAME);
         builder.field("query", query);
         builder.field("field", field);
+        if (vectorSimilarity != null) {
+            builder.field("similarity", vectorSimilarity);
+        }
         boostAndQueryNameToXContent(builder);
         builder.endObject();
     }
@@ -108,17 +117,17 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
             );
         }
         final DenseVectorFieldMapper.DenseVectorFieldType vectorFieldType = (DenseVectorFieldMapper.DenseVectorFieldType) fieldType;
-        return vectorFieldType.createExactKnnQuery(query);
+        return vectorFieldType.createExactKnnQuery(query, vectorSimilarity);
     }
 
     @Override
     protected boolean doEquals(ExactKnnQueryBuilder other) {
-        return field.equals(other.field) && Objects.equals(query, other.query);
+        return field.equals(other.field) && Objects.equals(query, other.query) && Objects.equals(vectorSimilarity, other.vectorSimilarity);
     }
 
     @Override
     protected int doHashCode() {
-        return Objects.hash(field, Objects.hashCode(query));
+        return Objects.hash(field, Objects.hashCode(query), vectorSimilarity);
     }
 
     @Override
diff --git a/server/src/main/java/org/elasticsearch/search/vectors/KnnScoreDocQueryBuilder.java b/server/src/main/java/org/elasticsearch/search/vectors/KnnScoreDocQueryBuilder.java
index 65f8c60297ad8..5cddc08b4b22d 100644
--- a/server/src/main/java/org/elasticsearch/search/vectors/KnnScoreDocQueryBuilder.java
+++ b/server/src/main/java/org/elasticsearch/search/vectors/KnnScoreDocQueryBuilder.java
@@ -37,6 +37,7 @@ public class KnnScoreDocQueryBuilder extends AbstractQueryBuilder<KnnScoreDocQue
     private final ScoreDoc[] scoreDocs;
     private final String fieldName;
     private final VectorData queryVector;
+    private final Float vectorSimilarity;
 
     /**
      * Creates a query builder.
@@ -44,20 +45,11 @@ public class KnnScoreDocQueryBuilder extends AbstractQueryBuilder<KnnScoreDocQue
      * @param scoreDocs the docs and scores this query should match. The array must be
      *                  sorted in order of ascending doc IDs.
      */
-    public KnnScoreDocQueryBuilder(ScoreDoc[] scoreDocs, String fieldName, float[] queryVector) {
-        this(scoreDocs, fieldName, VectorData.fromFloats(queryVector));
-    }
-
-    /**
-     * Creates a query builder.
-     *
-     * @param scoreDocs the docs and scores this query should match. The array must be
-     *                  sorted in order of ascending doc IDs.
-     */
-    public KnnScoreDocQueryBuilder(ScoreDoc[] scoreDocs, String fieldName, VectorData queryVector) {
+    public KnnScoreDocQueryBuilder(ScoreDoc[] scoreDocs, String fieldName, VectorData queryVector, Float vectorSimilarity) {
         this.scoreDocs = scoreDocs;
         this.fieldName = fieldName;
         this.queryVector = queryVector;
+        this.vectorSimilarity = vectorSimilarity;
     }
 
     public KnnScoreDocQueryBuilder(StreamInput in) throws IOException {
@@ -78,6 +70,11 @@ public KnnScoreDocQueryBuilder(StreamInput in) throws IOException {
             this.fieldName = null;
             this.queryVector = null;
         }
+        if (in.getTransportVersion().isPatchFrom(TransportVersions.FIX_VECTOR_SIMILARITY_INNER_HITS_BACKPORT_8_15)) {
+            this.vectorSimilarity = in.readOptionalFloat();
+        } else {
+            this.vectorSimilarity = null;
+        }
     }
 
     @Override
@@ -97,6 +94,10 @@ VectorData queryVector() {
         return queryVector;
     }
 
+    Float vectorSimilarity() {
+        return vectorSimilarity;
+    }
+
     @Override
     protected void doWriteTo(StreamOutput out) throws IOException {
         out.writeArray(Lucene::writeScoreDoc, scoreDocs);
@@ -113,6 +114,9 @@ protected void doWriteTo(StreamOutput out) throws IOException {
                 out.writeBoolean(false);
             }
         }
+        if (out.getTransportVersion().isPatchFrom(TransportVersions.FIX_VECTOR_SIMILARITY_INNER_HITS_BACKPORT_8_15)) {
+            out.writeOptionalFloat(vectorSimilarity);
+        }
     }
 
     @Override
@@ -129,6 +133,9 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
         if (queryVector != null) {
             builder.field("query", queryVector);
         }
+        if (vectorSimilarity != null) {
+            builder.field("similarity", vectorSimilarity);
+        }
         boostAndQueryNameToXContent(builder);
         builder.endObject();
     }
@@ -154,7 +161,7 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws
             return new MatchNoneQueryBuilder("The \"" + getName() + "\" query was rewritten to a \"match_none\" query.");
         }
         if (queryRewriteContext.convertToInnerHitsRewriteContext() != null && queryVector != null && fieldName != null) {
-            return new ExactKnnQueryBuilder(queryVector, fieldName);
+            return new ExactKnnQueryBuilder(queryVector, fieldName, vectorSimilarity);
         }
         return super.doRewrite(queryRewriteContext);
     }
@@ -193,7 +200,9 @@ protected boolean doEquals(KnnScoreDocQueryBuilder other) {
                 return false;
             }
         }
-        return Objects.equals(fieldName, other.fieldName) && Objects.equals(queryVector, other.queryVector);
+        return Objects.equals(fieldName, other.fieldName)
+            && Objects.equals(queryVector, other.queryVector)
+            && Objects.equals(vectorSimilarity, other.vectorSimilarity);
     }
 
     @Override
@@ -203,7 +212,7 @@ protected int doHashCode() {
             int hashCode = Objects.hash(scoreDoc.doc, scoreDoc.score, scoreDoc.shardIndex);
             result = 31 * result + hashCode;
         }
-        return Objects.hash(result, fieldName, Objects.hashCode(queryVector));
+        return Objects.hash(result, fieldName, vectorSimilarity, Objects.hashCode(queryVector));
     }
 
     @Override
diff --git a/server/src/main/java/org/elasticsearch/search/vectors/KnnSearchBuilder.java b/server/src/main/java/org/elasticsearch/search/vectors/KnnSearchBuilder.java
index 601c55293418d..519472f812a66 100644
--- a/server/src/main/java/org/elasticsearch/search/vectors/KnnSearchBuilder.java
+++ b/server/src/main/java/org/elasticsearch/search/vectors/KnnSearchBuilder.java
@@ -403,6 +403,10 @@ public KnnVectorQueryBuilder toQueryBuilder() {
             .addFilterQueries(filterQueries);
     }
 
+    public Float getSimilarity() {
+        return similarity;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
diff --git a/server/src/main/java/org/elasticsearch/search/vectors/KnnVectorQueryBuilder.java b/server/src/main/java/org/elasticsearch/search/vectors/KnnVectorQueryBuilder.java
index 0f64859e877f4..60bd80cfc4d97 100644
--- a/server/src/main/java/org/elasticsearch/search/vectors/KnnVectorQueryBuilder.java
+++ b/server/src/main/java/org/elasticsearch/search/vectors/KnnVectorQueryBuilder.java
@@ -400,7 +400,7 @@ protected QueryBuilder doRewrite(QueryRewriteContext ctx) throws IOException {
             ).queryName(queryName).addFilterQueries(filterQueries);
         }
         if (ctx.convertToInnerHitsRewriteContext() != null) {
-            return new ExactKnnQueryBuilder(queryVector, fieldName).boost(boost).queryName(queryName);
+            return new ExactKnnQueryBuilder(queryVector, fieldName, vectorSimilarity).boost(boost).queryName(queryName);
         }
         boolean changed = false;
         List<QueryBuilder> rewrittenQueries = new ArrayList<>(filterQueries.size());
diff --git a/server/src/test/java/org/elasticsearch/action/search/DfsQueryPhaseTests.java b/server/src/test/java/org/elasticsearch/action/search/DfsQueryPhaseTests.java
index e9ff8336ef4c9..6df4c379c5c8f 100644
--- a/server/src/test/java/org/elasticsearch/action/search/DfsQueryPhaseTests.java
+++ b/server/src/test/java/org/elasticsearch/action/search/DfsQueryPhaseTests.java
@@ -351,12 +351,14 @@ public void testRewriteShardSearchRequestWithRank() {
         KnnScoreDocQueryBuilder ksdqb0 = new KnnScoreDocQueryBuilder(
             new ScoreDoc[] { new ScoreDoc(1, 3.0f, 1), new ScoreDoc(4, 1.5f, 1) },
             "vector",
-            new float[] { 0.0f }
+            VectorData.fromFloats(new float[] { 0.0f }),
+            null
         );
         KnnScoreDocQueryBuilder ksdqb1 = new KnnScoreDocQueryBuilder(
             new ScoreDoc[] { new ScoreDoc(1, 2.0f, 1) },
             "vector2",
-            new float[] { 0.0f }
+            VectorData.fromFloats(new float[] { 0.0f }),
+            null
         );
         assertEquals(
             List.of(bm25, ksdqb0, ksdqb1),
diff --git a/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java b/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java
index 2a4554091dc91..27dbc79333c2d 100644
--- a/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java
+++ b/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java
@@ -215,7 +215,7 @@ public void testExactKnnQuery() {
             for (int i = 0; i < dims; i++) {
                 queryVector[i] = randomFloat();
             }
-            Query query = field.createExactKnnQuery(VectorData.fromFloats(queryVector));
+            Query query = field.createExactKnnQuery(VectorData.fromFloats(queryVector), null);
             assertTrue(query instanceof DenseVectorQuery.Floats);
         }
         {
@@ -233,7 +233,7 @@ public void testExactKnnQuery() {
             for (int i = 0; i < dims; i++) {
                 queryVector[i] = randomByte();
             }
-            Query query = field.createExactKnnQuery(VectorData.fromBytes(queryVector));
+            Query query = field.createExactKnnQuery(VectorData.fromBytes(queryVector), null);
             assertTrue(query instanceof DenseVectorQuery.Bytes);
         }
     }
diff --git a/server/src/test/java/org/elasticsearch/search/vectors/AbstractKnnVectorQueryBuilderTestCase.java b/server/src/test/java/org/elasticsearch/search/vectors/AbstractKnnVectorQueryBuilderTestCase.java
index f5d9f35e34695..0cc151a16e4b7 100644
--- a/server/src/test/java/org/elasticsearch/search/vectors/AbstractKnnVectorQueryBuilderTestCase.java
+++ b/server/src/test/java/org/elasticsearch/search/vectors/AbstractKnnVectorQueryBuilderTestCase.java
@@ -21,6 +21,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
+import org.elasticsearch.index.query.InnerHitsRewriteContext;
 import org.elasticsearch.index.query.MatchNoneQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryBuilders;
@@ -296,6 +297,22 @@ private void assertBWCSerialization(QueryBuilder newQuery, QueryBuilder bwcQuery
         }
     }
 
+    public void testRewriteForInnerHits() throws IOException {
+        SearchExecutionContext context = createSearchExecutionContext();
+        InnerHitsRewriteContext innerHitsRewriteContext = new InnerHitsRewriteContext(context.getParserConfig(), System::currentTimeMillis);
+        KnnVectorQueryBuilder queryBuilder = createTestQueryBuilder();
+        queryBuilder.boost(randomFloat());
+        queryBuilder.queryName(randomAlphaOfLength(10));
+        QueryBuilder rewritten = queryBuilder.rewrite(innerHitsRewriteContext);
+        assertTrue(rewritten instanceof ExactKnnQueryBuilder);
+        ExactKnnQueryBuilder exactKnnQueryBuilder = (ExactKnnQueryBuilder) rewritten;
+        assertEquals(queryBuilder.queryVector(), exactKnnQueryBuilder.getQuery());
+        assertEquals(queryBuilder.getFieldName(), exactKnnQueryBuilder.getField());
+        assertEquals(queryBuilder.boost(), exactKnnQueryBuilder.boost(), 0.0001f);
+        assertEquals(queryBuilder.queryName(), exactKnnQueryBuilder.queryName());
+        assertEquals(queryBuilder.getVectorSimilarity(), exactKnnQueryBuilder.vectorSimilarity());
+    }
+
     public void testRewriteWithQueryVectorBuilder() throws Exception {
         int dims = randomInt(1024);
         float[] expectedArray = new float[dims];
diff --git a/server/src/test/java/org/elasticsearch/search/vectors/ExactKnnQueryBuilderTests.java b/server/src/test/java/org/elasticsearch/search/vectors/ExactKnnQueryBuilderTests.java
index 627f8a184a147..220682f11de01 100644
--- a/server/src/test/java/org/elasticsearch/search/vectors/ExactKnnQueryBuilderTests.java
+++ b/server/src/test/java/org/elasticsearch/search/vectors/ExactKnnQueryBuilderTests.java
@@ -62,12 +62,12 @@ protected ExactKnnQueryBuilder doCreateTestQueryBuilder() {
         for (int i = 0; i < VECTOR_DIMENSION; i++) {
             query[i] = randomFloat();
         }
-        return new ExactKnnQueryBuilder(query, VECTOR_FIELD);
+        return new ExactKnnQueryBuilder(VectorData.fromFloats(query), VECTOR_FIELD, randomBoolean() ? randomFloat() : null);
     }
 
     @Override
     public void testValidOutput() {
-        ExactKnnQueryBuilder query = new ExactKnnQueryBuilder(new float[] { 1.0f, 2.0f, 3.0f }, "field");
+        ExactKnnQueryBuilder query = new ExactKnnQueryBuilder(VectorData.fromFloats(new float[] { 1.0f, 2.0f, 3.0f }), "field", null);
         String expected = """
             {
               "exact_knn" : {
@@ -80,10 +80,29 @@ public void testValidOutput() {
               }
             }""";
         assertEquals(expected, query.toString());
+        query = new ExactKnnQueryBuilder(VectorData.fromFloats(new float[] { 1.0f, 2.0f, 3.0f }), "field", 1f);
+        expected = """
+            {
+              "exact_knn" : {
+                "query" : [
+                  1.0,
+                  2.0,
+                  3.0
+                ],
+                "field" : "field",
+                "similarity" : 1.0
+              }
+            }""";
+        assertEquals(expected, query.toString());
     }
 
     @Override
     protected void doAssertLuceneQuery(ExactKnnQueryBuilder queryBuilder, Query query, SearchExecutionContext context) throws IOException {
+        if (queryBuilder.vectorSimilarity() != null) {
+            assertTrue(query instanceof VectorSimilarityQuery);
+            VectorSimilarityQuery vectorSimilarityQuery = (VectorSimilarityQuery) query;
+            query = vectorSimilarityQuery.getInnerKnnQuery();
+        }
         assertTrue(query instanceof DenseVectorQuery.Floats);
         DenseVectorQuery.Floats denseVectorQuery = (DenseVectorQuery.Floats) query;
         assertEquals(VECTOR_FIELD, denseVectorQuery.field);
diff --git a/server/src/test/java/org/elasticsearch/search/vectors/KnnScoreDocQueryBuilderTests.java b/server/src/test/java/org/elasticsearch/search/vectors/KnnScoreDocQueryBuilderTests.java
index 67bc6bde9c1af..3733c884c401b 100644
--- a/server/src/test/java/org/elasticsearch/search/vectors/KnnScoreDocQueryBuilderTests.java
+++ b/server/src/test/java/org/elasticsearch/search/vectors/KnnScoreDocQueryBuilderTests.java
@@ -64,7 +64,8 @@ protected KnnScoreDocQueryBuilder doCreateTestQueryBuilder() {
         return new KnnScoreDocQueryBuilder(
             scoreDocs.toArray(new ScoreDoc[0]),
             randomBoolean() ? "field" : null,
-            randomBoolean() ? randomVector(10) : null
+            randomBoolean() ? VectorData.fromFloats(randomVector(10)) : null,
+            randomBoolean() ? randomFloat() : null
         );
     }
 
@@ -73,7 +74,8 @@ public void testValidOutput() {
         KnnScoreDocQueryBuilder query = new KnnScoreDocQueryBuilder(
             new ScoreDoc[] { new ScoreDoc(0, 4.25f), new ScoreDoc(5, 1.6f) },
             "field",
-            new float[] { 1.0f, 2.0f }
+            VectorData.fromFloats(new float[] { 1.0f, 2.0f }),
+            null
         );
         String expected = """
             {
@@ -163,7 +165,8 @@ public void testRewriteToMatchNone() throws IOException {
         KnnScoreDocQueryBuilder queryBuilder = new KnnScoreDocQueryBuilder(
             new ScoreDoc[0],
             randomBoolean() ? "field" : null,
-            randomBoolean() ? randomVector(10) : null
+            randomBoolean() ? VectorData.fromFloats(randomVector(10)) : null,
+            randomBoolean() ? randomFloat() : null
         );
         QueryRewriteContext context = randomBoolean()
             ? new InnerHitsRewriteContext(createSearchExecutionContext().getParserConfig(), System::currentTimeMillis)
@@ -177,7 +180,8 @@ public void testRewriteForInnerHits() throws IOException {
         KnnScoreDocQueryBuilder queryBuilder = new KnnScoreDocQueryBuilder(
             new ScoreDoc[] { new ScoreDoc(0, 4.25f), new ScoreDoc(5, 1.6f) },
             randomAlphaOfLength(10),
-            randomVector(10)
+            VectorData.fromFloats(randomVector(10)),
+            randomBoolean() ? randomFloat() : null
         );
         queryBuilder.boost(randomFloat());
         queryBuilder.queryName(randomAlphaOfLength(10));
@@ -188,6 +192,7 @@ public void testRewriteForInnerHits() throws IOException {
         assertEquals(queryBuilder.fieldName(), exactKnnQueryBuilder.getField());
         assertEquals(queryBuilder.boost(), exactKnnQueryBuilder.boost(), 0.0001f);
         assertEquals(queryBuilder.queryName(), exactKnnQueryBuilder.queryName());
+        assertEquals(queryBuilder.vectorSimilarity(), exactKnnQueryBuilder.vectorSimilarity());
     }
 
     @Override
@@ -226,7 +231,12 @@ public void testScoreDocQueryWeightCount() throws IOException {
                 }
                 ScoreDoc[] scoreDocs = scoreDocsList.toArray(new ScoreDoc[0]);
 
-                KnnScoreDocQueryBuilder queryBuilder = new KnnScoreDocQueryBuilder(scoreDocs, "field", randomVector(10));
+                KnnScoreDocQueryBuilder queryBuilder = new KnnScoreDocQueryBuilder(
+                    scoreDocs,
+                    "field",
+                    VectorData.fromFloats(randomVector(10)),
+                    null
+                );
                 Query query = queryBuilder.doToQuery(context);
                 final Weight w = query.createWeight(searcher, ScoreMode.TOP_SCORES, 1.0f);
                 for (LeafReaderContext leafReaderContext : searcher.getLeafContexts()) {
@@ -269,7 +279,12 @@ public void testScoreDocQuery() throws IOException {
                 }
                 ScoreDoc[] scoreDocs = scoreDocsList.toArray(new ScoreDoc[0]);
 
-                KnnScoreDocQueryBuilder queryBuilder = new KnnScoreDocQueryBuilder(scoreDocs, "field", randomVector(10));
+                KnnScoreDocQueryBuilder queryBuilder = new KnnScoreDocQueryBuilder(
+                    scoreDocs,
+                    "field",
+                    VectorData.fromFloats(randomVector(10)),
+                    null
+                );
                 final Query query = queryBuilder.doToQuery(context);
                 final Weight w = query.createWeight(searcher, ScoreMode.TOP_SCORES, 1.0f);
 

From bbc2877d8051971bc94417e191a1fbcc3655a3f1 Mon Sep 17 00:00:00 2001
From: Benjamin Trent <4357155+benwtrent@users.noreply.github.com>
Date: Mon, 29 Jul 2024 16:24:22 -0400
Subject: [PATCH 2/4] fixing for backport

---
 .../test/search.vectors/100_knn_nested_search.yml             | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/100_knn_nested_search.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/100_knn_nested_search.yml
index d627be2fb15c3..d255a644183dc 100644
--- a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/100_knn_nested_search.yml
+++ b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/100_knn_nested_search.yml
@@ -414,8 +414,8 @@ setup:
 ---
 "nested Knn search with required similarity appropriately filters inner_hits":
   - requires:
-      cluster_features: "gte_v8.16.0"
-      reason: 'bugfix for 8.16'
+      cluster_features: "gte_v8.15.0"
+      reason: 'bugfix for 8.15'
 
   - do:
       search:

From 207a1c9b74ba000341b88074cc0ca5509fdf1ac4 Mon Sep 17 00:00:00 2001
From: Benjamin Trent <4357155+benwtrent@users.noreply.github.com>
Date: Mon, 29 Jul 2024 16:27:18 -0400
Subject: [PATCH 3/4] adj for backport

---
 .../src/main/java/org/elasticsearch/TransportVersions.java  | 1 -
 .../elasticsearch/search/vectors/ExactKnnQueryBuilder.java  | 6 ++----
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/TransportVersions.java b/server/src/main/java/org/elasticsearch/TransportVersions.java
index f2da6aba50da9..9d57c50b17aaf 100644
--- a/server/src/main/java/org/elasticsearch/TransportVersions.java
+++ b/server/src/main/java/org/elasticsearch/TransportVersions.java
@@ -211,7 +211,6 @@ static TransportVersion def(int id) {
     public static final TransportVersion ML_INFERENCE_AMAZON_BEDROCK_ADDED = def(8_702_00_0);
     public static final TransportVersion ENTERPRISE_GEOIP_DOWNLOADER_BACKPORT_8_15 = def(8_702_00_1);
     public static final TransportVersion FIX_VECTOR_SIMILARITY_INNER_HITS_BACKPORT_8_15 = def(8_702_00_2);
-    public static final TransportVersion FIX_VECTOR_SIMILARITY_INNER_HITS = def(8_713_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,
diff --git a/server/src/main/java/org/elasticsearch/search/vectors/ExactKnnQueryBuilder.java b/server/src/main/java/org/elasticsearch/search/vectors/ExactKnnQueryBuilder.java
index ebefd7d8b85b6..7e23075b1932e 100644
--- a/server/src/main/java/org/elasticsearch/search/vectors/ExactKnnQueryBuilder.java
+++ b/server/src/main/java/org/elasticsearch/search/vectors/ExactKnnQueryBuilder.java
@@ -54,8 +54,7 @@ public ExactKnnQueryBuilder(StreamInput in) throws IOException {
             this.query = VectorData.fromFloats(in.readFloatArray());
         }
         this.field = in.readString();
-        if (in.getTransportVersion().onOrAfter(TransportVersions.FIX_VECTOR_SIMILARITY_INNER_HITS)
-            || in.getTransportVersion().isPatchFrom(TransportVersions.FIX_VECTOR_SIMILARITY_INNER_HITS_BACKPORT_8_15)) {
+        if (in.getTransportVersion().isPatchFrom(TransportVersions.FIX_VECTOR_SIMILARITY_INNER_HITS_BACKPORT_8_15)) {
             this.vectorSimilarity = in.readOptionalFloat();
         } else {
             this.vectorSimilarity = null;
@@ -87,8 +86,7 @@ protected void doWriteTo(StreamOutput out) throws IOException {
             out.writeFloatArray(query.asFloatVector());
         }
         out.writeString(field);
-        if (out.getTransportVersion().onOrAfter(TransportVersions.FIX_VECTOR_SIMILARITY_INNER_HITS)
-            || out.getTransportVersion().isPatchFrom(TransportVersions.FIX_VECTOR_SIMILARITY_INNER_HITS_BACKPORT_8_15)) {
+        if (out.getTransportVersion().isPatchFrom(TransportVersions.FIX_VECTOR_SIMILARITY_INNER_HITS_BACKPORT_8_15)) {
             out.writeOptionalFloat(vectorSimilarity);
         }
     }

From b5a4972dc379085b94010e51d296aa3efc814a73 Mon Sep 17 00:00:00 2001
From: Benjamin Trent <4357155+benwtrent@users.noreply.github.com>
Date: Mon, 29 Jul 2024 17:09:16 -0400
Subject: [PATCH 4/4] fix compilation for tests

---
 .../java/org/elasticsearch/action/search/DfsQueryPhaseTests.java | 1 +
 1 file changed, 1 insertion(+)

diff --git a/server/src/test/java/org/elasticsearch/action/search/DfsQueryPhaseTests.java b/server/src/test/java/org/elasticsearch/action/search/DfsQueryPhaseTests.java
index 6df4c379c5c8f..511a7c90dbedc 100644
--- a/server/src/test/java/org/elasticsearch/action/search/DfsQueryPhaseTests.java
+++ b/server/src/test/java/org/elasticsearch/action/search/DfsQueryPhaseTests.java
@@ -35,6 +35,7 @@
 import org.elasticsearch.search.rank.TestRankBuilder;
 import org.elasticsearch.search.vectors.KnnScoreDocQueryBuilder;
 import org.elasticsearch.search.vectors.KnnSearchBuilder;
+import org.elasticsearch.search.vectors.VectorData;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.InternalAggregationTestCase;
 import org.elasticsearch.transport.Transport;