Improve explanation in rescore

rest-api-spec/src/main/resources/rest-api-spec/test/search/210_rescore_explain.yml

@@ -0,0 +1,39 @@
+---
+"Score should match explanation in rescore":
+  - skip:
+      version: " - 6.99.99"
+      reason: Explanation for rescoring was corrected after these versions
+  - do:
+      bulk:
+        refresh: true
+        body:
+          - '{"index": {"_index": "test_index", "_type": "_doc", "_id": "1"}}'
+          - '{"f1": "1"}'
+          - '{"index": {"_index": "test_index", "_type": "_doc", "_id": "2"}}'
+          - '{"f1": "2"}'
+          - '{"index": {"_index": "test_index", "_type": "_doc", "_id": "3"}}'
+          - '{"f1": "3"}'
+
+  - do:
+      search:
+        index: test_index
+        body:
+          explain: true
+          query:
+            match_all: {}
+          rescore:
+            window_size: 2
+            query:
+              rescore_query:
+                match_all: {}
+              query_weight: 5
+              rescore_query_weight: 10
+
+  - match: { hits.hits.0._score: 15 }
+  - match: { hits.hits.0._explanation.value: 15 }
+
+  - match: { hits.hits.1._score: 15 }
+  - match: { hits.hits.1._explanation.value: 15 }
+
+  - match: { hits.hits.2._score: 5 }
+  - match: { hits.hits.2._explanation.value: 5 }

server/src/main/java/org/elasticsearch/search/rescore/QueryRescorer.java

@@ -30,6 +30,8 @@
 import java.util.Arrays;
 import java.util.Comparator;
 import java.util.Set;
+import java.util.Collections;
+import static java.util.stream.Collectors.toSet;
 
 public final class QueryRescorer implements Rescorer {
 
@@ -61,6 +63,11 @@ protected float combine(float firstPassScore, boolean secondPassMatches, float s
         // First take top slice of incoming docs, to be rescored:
         TopDocs topNFirstPass = topN(topDocs, rescoreContext.getWindowSize());
 
+        // Save doc IDs for which rescoring was applied to be used in score explanation
+        Set<Integer> topNDocIDs = Collections.unmodifiableSet(
+            Arrays.stream(topNFirstPass.scoreDocs).map(scoreDoc -> scoreDoc.doc).collect(toSet()));
+        rescoreContext.setRescoredDocs(topNDocIDs);
+
         // Rescore them:
         TopDocs rescored = rescorer.rescore(searcher, topNFirstPass, rescoreContext.getWindowSize());
 
@@ -76,8 +83,6 @@ public Explanation explain(int topLevelDocId, IndexSearcher searcher, RescoreCon
             // this should not happen but just in case
             return Explanation.noMatch("nothing matched");
         }
-        // TODO: this isn't right?  I.e., we are incorrectly pretending all first pass hits were rescored?  If the requested docID was
-        // beyond the top rescoreContext.window() in the first pass hits, we don't rescore it now?
         Explanation rescoreExplain = searcher.explain(rescore.query(), topLevelDocId);
         float primaryWeight = rescore.queryWeight();
 
@@ -92,7 +97,7 @@ public Explanation explain(int topLevelDocId, IndexSearcher searcher, RescoreCon
 
         // NOTE: we don't use Lucene's Rescorer.explain because we want to insert our own description with which ScoreMode was used.  Maybe
         // we should add QueryRescorer.explainCombine to Lucene?
-        if (rescoreExplain != null && rescoreExplain.isMatch()) {
+        if (rescoreContext.isRescored(topLevelDocId) && rescoreExplain != null && rescoreExplain.isMatch()) {
             float secondaryWeight = rescore.rescoreQueryWeight();
             Explanation sec = Explanation.match(
                     rescoreExplain.getValue() * secondaryWeight,

server/src/main/java/org/elasticsearch/search/rescore/RescoreContext.java

@@ -19,6 +19,8 @@
 
 package org.elasticsearch.search.rescore;
 
+import java.util.Set;
+
 /**
  * Context available to the rescore while it is running. Rescore
  * implementations should extend this with any additional resources that
@@ -27,6 +29,7 @@
 public class RescoreContext {
     private final int windowSize;
     private final Rescorer rescorer;
+    private Set<Integer> recroredDocs; //doc Ids for which rescoring was applied
 
     /**
      * Build the context.
@@ -50,4 +53,12 @@ public Rescorer rescorer() {
     public int getWindowSize() {
         return windowSize;
     }
+
+    public void setRescoredDocs(Set<Integer> docIds) {
+        recroredDocs = docIds;
+    }
+
+    public boolean isRescored(int docId) {
+        return recroredDocs.contains(docId);
+    }
 }