From 7ca753b2cb1235d9eb6847d5ab43693447007c5f Mon Sep 17 00:00:00 2001
From: Jim Ferenczi <jim.ferenczi@elastic.co>
Date: Fri, 21 Jun 2024 12:13:57 +0100
Subject: [PATCH] minor fixes

---
 .../elasticsearch/search/retriever/RankDocsQueryBuilder.java   | 2 ++
 .../search/retriever/RankDocsRetrieverBuilder.java             | 1 +
 .../org/elasticsearch/search/retriever/RetrieverBuilder.java   | 3 ++-
 3 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/server/src/main/java/org/elasticsearch/search/retriever/RankDocsQueryBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/RankDocsQueryBuilder.java
index 699d6e5c8eb2b..52546f9af223e 100644
--- a/server/src/main/java/org/elasticsearch/search/retriever/RankDocsQueryBuilder.java
+++ b/server/src/main/java/org/elasticsearch/search/retriever/RankDocsQueryBuilder.java
@@ -78,6 +78,7 @@ private static int[] findSegmentStarts(IndexReader reader, RankDoc[] docs) {
 
     @Override
     protected void doXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject(NAME);
         builder.startArray("docs");
         for (RankDoc doc : rankDocs) {
             builder.startObject();
@@ -85,6 +86,7 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
             builder.endObject();
         }
         builder.endArray();
+        builder.endObject();
     }
 
     @Override
diff --git a/server/src/main/java/org/elasticsearch/search/retriever/RankDocsRetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/RankDocsRetrieverBuilder.java
index 00f101ca9f589..cdd389b98f648 100644
--- a/server/src/main/java/org/elasticsearch/search/retriever/RankDocsRetrieverBuilder.java
+++ b/server/src/main/java/org/elasticsearch/search/retriever/RankDocsRetrieverBuilder.java
@@ -10,6 +10,7 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.index.query.BoolQueryBuilder;
 import org.elasticsearch.index.query.DisMaxQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
diff --git a/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java
index 60806079c09c0..43057b0ddd126 100644
--- a/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java
+++ b/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java
@@ -203,7 +203,8 @@ public RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOException {
     public abstract QueryBuilder topDocsQuery();
 
     /**
-     * Called by {@link CombineRetrieverBuilder} to seal the top docs returned by this retriever.
+     * This function is called by {@link CombineRetrieverBuilder}s to seal the top docs returned by
+     * this retriever during rewrite.
      */
     public void setRankDocs(RankDoc[] rankDocs) {
         this.rankDocs = rankDocs;