From 5a17826fc22c76bdd97f98ce63bad766e83697f6 Mon Sep 17 00:00:00 2001
From: Jim Ferenczi <jim.ferenczi@elastic.co>
Date: Mon, 14 Oct 2024 13:41:15 +0100
Subject: [PATCH] Fix standard retriever rewrite

Closes #114466
---
 .../retriever/rankdoc/RankDocsQuery.java      |  9 ++++-----
 .../rankdoc/RankDocsQueryBuilderTests.java    | 20 +++++++++++++++++++
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/search/retriever/rankdoc/RankDocsQuery.java b/server/src/main/java/org/elasticsearch/search/retriever/rankdoc/RankDocsQuery.java
index fb5015a82dbdb..b78d9e40ba120 100644
--- a/server/src/main/java/org/elasticsearch/search/retriever/rankdoc/RankDocsQuery.java
+++ b/server/src/main/java/org/elasticsearch/search/retriever/rankdoc/RankDocsQuery.java
@@ -107,11 +107,10 @@ public Explanation explain(LeafReaderContext context, int doc) throws IOExceptio
 
                 @Override
                 public Scorer scorer(LeafReaderContext context) {
-                    // Segment starts indicate how many docs are in the segment,
-                    // upper equalling lower indicates no documents for this segment
-                    if (segmentStarts[context.ord] == segmentStarts[context.ord + 1]) {
-                        return null;
-                    }
+                    /**
+                     * We return a scorer even if there are no ranked documents within the segment.
+                     * This ensures the correct propagation of the maximum score.
+                     */
                     return new Scorer(this) {
                         final int lower = segmentStarts[context.ord];
                         final int upper = segmentStarts[context.ord + 1];
diff --git a/server/src/test/java/org/elasticsearch/search/retriever/rankdoc/RankDocsQueryBuilderTests.java b/server/src/test/java/org/elasticsearch/search/retriever/rankdoc/RankDocsQueryBuilderTests.java
index ca05c57b7d733..f1aee5764fba6 100644
--- a/server/src/test/java/org/elasticsearch/search/retriever/rankdoc/RankDocsQueryBuilderTests.java
+++ b/server/src/test/java/org/elasticsearch/search/retriever/rankdoc/RankDocsQueryBuilderTests.java
@@ -195,6 +195,26 @@ public void testRankDocsQueryEarlyTerminate() throws IOException {
                     assertThat(col.totalHits.value, equalTo((long) topSize));
                     assertEqualTopDocs(col.scoreDocs, rankDocs);
                 }
+
+                {
+                    // A single rank doc in the last segment
+                    RankDoc[] singleRankDoc = new RankDoc[1];
+                    singleRankDoc[0] = rankDocs[rankDocs.length-1];
+                    RankDocsQuery q = new RankDocsQuery(
+                            reader,
+                            singleRankDoc,
+                            new Query[] { NumericDocValuesField.newSlowExactQuery("active", 1) },
+                            new String[1],
+                            false
+                    );
+                    var topDocsManager = new TopScoreDocCollectorManager(1, null, 0);
+                    var col = searcher.search(q, topDocsManager);
+                    assertThat(
+                            col.totalHits.value,
+                            lessThanOrEqualTo((long) (2 + rankDocs.length))
+                    );
+                    assertEqualTopDocs(col.scoreDocs, singleRankDoc);
+                }
             }
         }
     }