Fix exception when merging completion suggestions (#70414)

Under certain circumstances we can get "array_index_out_of_bounds" exceptions in
the fetch phase when merging comletion suggestion results with additonal field
collapsing in place. We store the ScoreDoc array in the SortedTopDocs score docs
contain both regular search hits and completion suggestion results added in
SearchPhaseController#sortDocs, so when merging these we need to know the
correct array offset where the completion results begin. This is based on the
hits length calculated in SearchPhaseController#getHits, which doesn't take into
account that there might be suggestion results present when calculating the
number of hits. This change adds the number of suggestions to SortedTopDocs in
order to be able to later account for it when calculating the hits.

Closes #70328

Author: Christoph Büscher

server/src/internalClusterTest/java/org/elasticsearch/search/suggest/CompletionSuggestSearchIT.java

@@ -26,9 +26,11 @@
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentFactory;
 import org.elasticsearch.index.mapper.MapperParsingException;
+import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.search.aggregations.AggregationBuilders;
 import org.elasticsearch.search.aggregations.Aggregator.SubAggCollectionMode;
+import org.elasticsearch.search.collapse.CollapseBuilder;
 import org.elasticsearch.search.sort.FieldSortBuilder;
 import org.elasticsearch.search.suggest.completion.CompletionStats;
 import org.elasticsearch.search.suggest.completion.CompletionSuggestion;
@@ -1217,6 +1219,58 @@ public void testSuggestOnlyExplain() throws Exception {
         assertSuggestions(searchResponse, "foo", "suggestion10", "suggestion9", "suggestion8", "suggestion7", "suggestion6");
     }
 
+    public void testCompletionWithCollapse() throws Exception {
+        String suggestField = "suggest_field";
+        XContentBuilder mapping = jsonBuilder().startObject()
+            .startObject("properties")
+            .startObject("collapse_field")
+            .field("type", "keyword")
+            .endObject()
+            .startObject(suggestField)
+            .field("type", "completion")
+            .field("analyzer", "whitespace")
+            .endObject()
+            .endObject()
+            .endObject();
+
+        String index = "test";
+        assertAcked(
+            client().admin()
+                .indices()
+                .prepareCreate(index)
+                .setSettings(Settings.builder().put("index.number_of_shards", 2))
+                .setMapping(mapping)
+                .get()
+        );
+
+        int numDocs = 2;
+        for (int i = 0; i < numDocs; i++) {
+            XContentBuilder builder = jsonBuilder().startObject();
+            builder.startObject(suggestField).field("input", "suggestion" + i).field("weight", i).endObject();
+            builder.field("collapse_field", "collapse me").endObject();  // all docs the same value for collapsing
+            client().prepareIndex(index).setId("" + i).setSource(builder).get();
+        }
+        client().admin().indices().prepareRefresh(index).get();
+        CompletionSuggestionBuilder prefix = SuggestBuilders.completionSuggestion(suggestField).prefix("sug").size(1);
+
+        SearchResponse searchResponse = client().prepareSearch("test")
+            .setQuery(QueryBuilders.matchAllQuery())
+            .setFrom(1)
+            .setSize(1)
+            .setCollapse(new CollapseBuilder("collapse_field"))
+            .suggest(new SuggestBuilder().addSuggestion("the_suggestion", prefix))
+            .get();
+        assertAllSuccessful(searchResponse);
+
+        assertThat(searchResponse.getSuggest().getSuggestion("the_suggestion"), is(notNullValue()));
+        Suggest.Suggestion<Suggest.Suggestion.Entry<Suggest.Suggestion.Entry.Option>> suggestion = searchResponse.getSuggest()
+            .getSuggestion("the_suggestion");
+
+        List<String> suggestionList = getNames(suggestion.getEntries().get(0));
+        assertThat(suggestionList, contains("suggestion" + (numDocs - 1)));
+        assertEquals(0, searchResponse.getHits().getHits().length);
+    }
+
     public static boolean isReservedChar(char c) {
         switch (c) {
             case '\u001F':

server/src/main/java/org/elasticsearch/action/search/SearchPhaseController.java

@@ -149,8 +149,8 @@ static SortedTopDocs sortDocs(boolean ignoreFrom, final Collection<TopDocs> topD
         final TopDocs mergedTopDocs = mergeTopDocs(topDocs, size, ignoreFrom ? 0 : from);
         final ScoreDoc[] mergedScoreDocs = mergedTopDocs == null ? EMPTY_DOCS : mergedTopDocs.scoreDocs;
         ScoreDoc[] scoreDocs = mergedScoreDocs;
+        int numSuggestDocs = 0;
         if (reducedCompletionSuggestions.isEmpty() == false) {
-            int numSuggestDocs = 0;
             for (CompletionSuggestion completionSuggestion : reducedCompletionSuggestions) {
                 assert completionSuggestion != null;
                 numSuggestDocs += completionSuggestion.getOptions().size();
@@ -180,7 +180,7 @@ static SortedTopDocs sortDocs(boolean ignoreFrom, final Collection<TopDocs> topD
                 isSortedByField = true;
             }
         }
-        return new SortedTopDocs(scoreDocs, isSortedByField, sortFields, collapseField, collapseValues);
+        return new SortedTopDocs(scoreDocs, isSortedByField, sortFields, collapseField, collapseValues, numSuggestDocs);
     }
 
     static TopDocs mergeTopDocs(Collection<TopDocs> results, int topN, int from) {
@@ -316,7 +316,8 @@ private SearchHits getHits(ReducedQueryPhase reducedQueryPhase, boolean ignoreFr
         int from = ignoreFrom ? 0 : reducedQueryPhase.from;
         int numSearchHits = (int) Math.min(reducedQueryPhase.fetchHits - from, reducedQueryPhase.size);
         // with collapsing we can have more fetch hits than sorted docs
-        numSearchHits = Math.min(sortedTopDocs.scoreDocs.length, numSearchHits);
+        // also we need to take into account that we potentially have completion suggestions stored in the scoreDocs array
+        numSearchHits = Math.min(sortedTopDocs.scoreDocs.length - sortedTopDocs.numberOfCompletionsSuggestions, numSearchHits);
         // merge hits
         List<SearchHit> hits = new ArrayList<>();
         if (fetchResults.isEmpty() == false) {
@@ -665,7 +666,7 @@ void add(TopDocsAndMaxScore topDocs, boolean timedOut, Boolean terminatedEarly)
     }
 
     static final class SortedTopDocs {
-        static final SortedTopDocs EMPTY = new SortedTopDocs(EMPTY_DOCS, false, null, null, null);
+        static final SortedTopDocs EMPTY = new SortedTopDocs(EMPTY_DOCS, false, null, null, null, 0);
         // the searches merged top docs
         final ScoreDoc[] scoreDocs;
         // <code>true</code> iff the result score docs is sorted by a field (not score), this implies that <code>sortField</code> is set.
@@ -674,14 +675,16 @@ static final class SortedTopDocs {
         final SortField[] sortFields;
         final String collapseField;
         final Object[] collapseValues;
+        final int numberOfCompletionsSuggestions;
 
         SortedTopDocs(ScoreDoc[] scoreDocs, boolean isSortedByField, SortField[] sortFields,
-                      String collapseField, Object[] collapseValues) {
+                      String collapseField, Object[] collapseValues, int numberOfCompletionsSuggestions) {
             this.scoreDocs = scoreDocs;
             this.isSortedByField = isSortedByField;
             this.sortFields = sortFields;
             this.collapseField = collapseField;
             this.collapseValues = collapseValues;
+            this.numberOfCompletionsSuggestions = numberOfCompletionsSuggestions;
         }
     }
 }