Add test for knn retriever

carlosdelest · Nov 29, 2024 · 0dab8ea · 0dab8ea
1 parent 4fbbadd
commit 0dab8ea
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/server/src/main/java/org/elasticsearch/search/retriever/KnnRetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/KnnRetrieverBuilder.java
@@ -257,7 +257,11 @@ public void extractToSearchSourceBuilder(SearchSourceBuilder searchSourceBuilder
         searchSourceBuilder.knnSearch(knnSearchBuilders);
     }
 
-    // ---- FOR TESTING XCONTENT PARSING ----
+    RescoreVectorBuilder rescoreVectorBuilder() {
+        return rescoreVectorBuilder;
+    }
+
+// ---- FOR TESTING XCONTENT PARSING ----
 
     @Override
     public void doToXContent(XContentBuilder builder, Params params) throws IOException {
@@ -278,7 +282,9 @@ public void doToXContent(XContentBuilder builder, Params params) throws IOExcept
         }
 
         if (rescoreVectorBuilder != null) {
-            builder.field(RESCORE_FIELD.getPreferredName(), rescoreVectorBuilder);
+            builder.startObject(RESCORE_FIELD.getPreferredName());
+            rescoreVectorBuilder.toXContent(builder, params);
+            builder.endObject();
         }
     }
 

diff --git a/server/src/test/java/org/elasticsearch/search/retriever/KnnRetrieverBuilderParsingTests.java b/server/src/test/java/org/elasticsearch/search/retriever/KnnRetrieverBuilderParsingTests.java
@@ -105,6 +105,7 @@ public void testRewrite() throws IOException {
             assertNull(source.query());
             assertThat(source.knnSearch().size(), equalTo(1));
             assertThat(source.knnSearch().get(0).getFilterQueries().size(), equalTo(knnRetriever.preFilterQueryBuilders.size()));
+            assertThat(source.knnSearch().get(0).getRescoreVectorBuilder(), equalTo(knnRetriever.rescoreVectorBuilder()));
             for (int j = 0; j < knnRetriever.preFilterQueryBuilders.size(); j++) {
                 assertThat(
                     source.knnSearch().get(0).getFilterQueries().get(j),