apply review comments

rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search/112_field_collapsing_with_rescore.yml

@@ -0,0 +1,107 @@
+setup:
+  - skip:
+      version: " - 8.14.99"
+      reason: Collapse with rescore added in 8.15.0
+  - do:
+      indices.create:
+        index: products
+        body:
+          mappings:
+              properties:
+                product_id: { type: keyword }
+                description: { type: text }
+                popularity: { type: integer }
+
+  - do:
+      bulk:
+        index: products
+        refresh: true
+        body:
+          - '{"index": {"_id": "1", "routing": "0"}}'
+          - '{"product_id": "0", "description": "flat tv 4K HDR", "score": 2, "popularity": 30}'
+          - '{"index": {"_id": "2", "routing": "10"}}'
+          - '{"product_id": "10", "description": "LED Smart TV 32", "score": 5, "popularity": 100}'
+          - '{"index": {"_id": "3", "routing": "10"}}'
+          - '{"product_id": "10", "description": "LED Smart TV 65", "score": 10, "popularity": 50}'
+          - '{"index": {"_id": "4", "routing": "0"}}'
+          - '{"product_id": "0", "description": "flat tv", "score": 1, "popularity": 10}'
+          - '{"index": {"_id": "5", "routing": "129"}}'
+          - '{"product_id": "129", "description": "just a tv", "score": 100, "popularity": 3}'
+
+---
+"field collapsing and rescore":
+  - do:
+      search:
+        index: products
+        body:
+          query:
+            bool:
+              filter:
+                match:
+                  description: "tv"
+              should:
+                script_score:
+                  query: { match_all: { } }
+                  script:
+                    source: "doc['score'].value"
+          collapse:
+            field: product_id
+          rescore:
+            query:
+              rescore_query:
+                script_score:
+                  query: { match_all: { } }
+                  script:
+                    source: "doc['popularity'].value"
+              query_weight: 0
+              rescore_query_weight: 1
+
+
+  - match:  {hits.total.value: 5 }
+  - length: {hits.hits: 3 }
+  - match:  {hits.hits.0._id:  "3"}
+  - match:  {hits.hits.0._score:  50}
+  - match:  {hits.hits.0.fields.product_id:  ["10"]}
+  - match: { hits.hits.1._id: "1" }
+  - match: { hits.hits.1._score: 30 }
+  - match: { hits.hits.1.fields.product_id: ["0"] }
+  - match: { hits.hits.2._id: "5" }
+  - match: { hits.hits.2._score: 3 }
+  - match: { hits.hits.2.fields.product_id: ["129"] }
+
+---
+"field collapsing and rescore with window_size":
+  - do:
+      search:
+        index: products
+        body:
+          query:
+            bool:
+              filter:
+                match:
+                  description: "tv"
+              should:
+                script_score:
+                  query: { match_all: { } }
+                  script:
+                    source: "doc['score'].value"
+          collapse:
+            field: product_id
+          rescore:
+            window_size: 2
+            query:
+              rescore_query:
+                script_score:
+                  query: { match_all: { } }
+                  script:
+                    source: "doc['popularity'].value"
+              query_weight: 0
+              rescore_query_weight: 1
+          size: 1
+
+
+  - match:  {hits.total.value: 5 }
+  - length: {hits.hits: 1 }
+  - match:  {hits.hits.0._id:  "3"}
+  - match:  {hits.hits.0._score:  50}
+  - match:  {hits.hits.0.fields.product_id:  ["10"]}

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

@@ -14,11 +14,11 @@
 import org.apache.lucene.search.TopDocs;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.common.lucene.search.TopDocsAndMaxScore;
+import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.lucene.grouping.TopFieldGroups;
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
-import java.util.HashMap;
 import java.util.Map;
 
 /**
@@ -68,7 +68,7 @@ public static void execute(SearchContext context) {
     private static TopFieldGroups rewriteTopGroups(TopFieldGroups originalTopGroups, TopDocs rescoredTopDocs) {
         assert originalTopGroups.fields.length == 1 && SortField.FIELD_SCORE.equals(originalTopGroups.fields[0])
             : "rescore must always sort by score descending";
-        Map<Integer, Object> docIdToGroupValue = new HashMap<>();
+        Map<Integer, Object> docIdToGroupValue = Maps.newMapWithExpectedSize(originalTopGroups.scoreDocs.length);
         for (int i = 0; i < originalTopGroups.scoreDocs.length; i++) {
             docIdToGroupValue.put(originalTopGroups.scoreDocs[i].doc, originalTopGroups.groupValues[i]);
         }

server/src/test/java/org/elasticsearch/action/search/SearchRequestTests.java

@@ -17,7 +17,6 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.util.ArrayUtils;
 import org.elasticsearch.core.TimeValue;
-import org.elasticsearch.index.query.MatchAllQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.index.query.TermQueryBuilder;
 import org.elasticsearch.search.AbstractSearchTestCase;

server/src/test/java/org/elasticsearch/search/SearchServiceTests.java

@@ -100,7 +100,6 @@
 import org.elasticsearch.search.query.NonCountingTermQuery;
 import org.elasticsearch.search.query.QuerySearchRequest;
 import org.elasticsearch.search.query.QuerySearchResult;
-import org.elasticsearch.search.rescore.QueryRescorerBuilder;
 import org.elasticsearch.search.slice.SliceBuilder;
 import org.elasticsearch.search.suggest.SuggestBuilder;
 import org.elasticsearch.tasks.TaskCancelHelper;