Faster and shorter SearchPhaseController.reduceQueryPhase (#119855) (#122092)

We can avoid one list copy and some indirection by collecting to a list of non-null
query responses instead of non-null generic responses right away (this also avoids the unfortunate reassignment of
a method parameter).
Also, this method is fairly long, this at least removes all redundant local variables
and as a result a bit of computation in some cases.

Author: original-brownbear

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

@@ -420,9 +420,8 @@ static ReducedQueryPhase reducedQueryPhase(
         assert numReducePhases >= 0 : "num reduce phases must be >= 0 but was: " + numReducePhases;
         numReducePhases++; // increment for this phase
         if (queryResults.isEmpty()) { // early terminate we have nothing to reduce
-            final TotalHits totalHits = topDocsStats.getTotalHits();
             return new ReducedQueryPhase(
-                totalHits,
+                topDocsStats.getTotalHits(),
                 topDocsStats.fetchHits,
                 topDocsStats.getMaxScore(),
                 false,
@@ -439,8 +438,7 @@ static ReducedQueryPhase reducedQueryPhase(
                 true
             );
         }
-        int total = queryResults.size();
-        final Collection<SearchPhaseResult> nonNullResults = new ArrayList<>();
+        final List<QuerySearchResult> nonNullResults = new ArrayList<>();
         boolean hasSuggest = false;
         boolean hasProfileResults = false;
         for (SearchPhaseResult queryResult : queryResults) {
@@ -450,26 +448,24 @@ static ReducedQueryPhase reducedQueryPhase(
             }
             hasSuggest |= res.suggest() != null;
             hasProfileResults |= res.hasProfileResults();
-            nonNullResults.add(queryResult);
+            nonNullResults.add(res);
         }
-        queryResults = nonNullResults;
-        validateMergeSortValueFormats(queryResults);
-        if (queryResults.isEmpty()) {
-            var ex = new IllegalStateException("must have at least one non-empty search result, got 0 out of " + total);
+        validateMergeSortValueFormats(nonNullResults);
+        if (nonNullResults.isEmpty()) {
+            var ex = new IllegalStateException("must have at least one non-empty search result, got 0 out of " + queryResults.size());
             assert false : ex;
             throw ex;
         }
 
         // count the total (we use the query result provider here, since we might not get any hits (we scrolled past them))
         final Map<String, List<Suggestion<?>>> groupedSuggestions = hasSuggest ? new HashMap<>() : Collections.emptyMap();
         final Map<String, SearchProfileQueryPhaseResult> profileShardResults = hasProfileResults
-            ? Maps.newMapWithExpectedSize(queryResults.size())
+            ? Maps.newMapWithExpectedSize(nonNullResults.size())
             : Collections.emptyMap();
         int from = 0;
         int size = 0;
         DocValueFormat[] sortValueFormats = null;
-        for (SearchPhaseResult entry : queryResults) {
-            QuerySearchResult result = entry.queryResult();
+        for (QuerySearchResult result : nonNullResults) {
             from = result.from();
             // sorted queries can set the size to 0 if they have enough competitive hits.
             size = Math.max(result.size(), size);
@@ -480,49 +476,44 @@ static ReducedQueryPhase reducedQueryPhase(
             if (hasSuggest) {
                 assert result.suggest() != null;
                 for (Suggestion<? extends Suggestion.Entry<? extends Suggestion.Entry.Option>> suggestion : result.suggest()) {
-                    List<Suggestion<?>> suggestionList = groupedSuggestions.computeIfAbsent(suggestion.getName(), s -> new ArrayList<>());
-                    suggestionList.add(suggestion);
+                    groupedSuggestions.computeIfAbsent(suggestion.getName(), s -> new ArrayList<>()).add(suggestion);
                     if (suggestion instanceof CompletionSuggestion completionSuggestion) {
                         completionSuggestion.setShardIndex(result.getShardIndex());
                     }
                 }
             }
             assert bufferedTopDocs.isEmpty() || result.hasConsumedTopDocs() : "firstResult has no aggs but we got non null buffered aggs?";
             if (hasProfileResults) {
-                String key = result.getSearchShardTarget().toString();
-                profileShardResults.put(key, result.consumeProfileResult());
+                profileShardResults.put(result.getSearchShardTarget().toString(), result.consumeProfileResult());
             }
         }
-        final Suggest reducedSuggest;
-        final List<CompletionSuggestion> reducedCompletionSuggestions;
-        if (groupedSuggestions.isEmpty()) {
-            reducedSuggest = null;
-            reducedCompletionSuggestions = Collections.emptyList();
-        } else {
-            reducedSuggest = new Suggest(Suggest.reduce(groupedSuggestions));
-            reducedCompletionSuggestions = reducedSuggest.filter(CompletionSuggestion.class);
-        }
-        final SearchProfileResultsBuilder profileBuilder = profileShardResults.isEmpty()
-            ? null
-            : new SearchProfileResultsBuilder(profileShardResults);
+        final Suggest reducedSuggest = groupedSuggestions.isEmpty() ? null : new Suggest(Suggest.reduce(groupedSuggestions));
         final SortedTopDocs sortedTopDocs;
         if (queryPhaseRankCoordinatorContext == null) {
-            sortedTopDocs = sortDocs(isScrollRequest, bufferedTopDocs, from, size, reducedCompletionSuggestions);
+            sortedTopDocs = sortDocs(
+                isScrollRequest,
+                bufferedTopDocs,
+                from,
+                size,
+                reducedSuggest == null ? Collections.emptyList() : reducedSuggest.filter(CompletionSuggestion.class)
+            );
         } else {
-            ScoreDoc[] rankedDocs = queryPhaseRankCoordinatorContext.rankQueryPhaseResults(
-                queryResults.stream().map(SearchPhaseResult::queryResult).toList(),
-                topDocsStats
+            sortedTopDocs = new SortedTopDocs(
+                queryPhaseRankCoordinatorContext.rankQueryPhaseResults(nonNullResults, topDocsStats),
+                false,
+                null,
+                null,
+                null,
+                0
             );
-            sortedTopDocs = new SortedTopDocs(rankedDocs, false, null, null, null, 0);
             size = sortedTopDocs.scoreDocs.length;
             // we need to reset from here as pagination and result trimming has already taken place
             // within the `QueryPhaseRankCoordinatorContext#rankQueryPhaseResults` and we don't want
             // to apply it again in the `getHits` method.
             from = 0;
         }
-        final TotalHits totalHits = topDocsStats.getTotalHits();
         return new ReducedQueryPhase(
-            totalHits,
+            topDocsStats.getTotalHits(),
             topDocsStats.fetchHits,
             topDocsStats.getMaxScore(),
             topDocsStats.timedOut,
@@ -534,7 +525,7 @@ static ReducedQueryPhase reducedQueryPhase(
                     bufferedAggs,
                     performFinalReduce ? aggReduceContextBuilder.forFinalReduction() : aggReduceContextBuilder.forPartialReduction()
                 ),
-            profileBuilder,
+            profileShardResults.isEmpty() ? null : new SearchProfileResultsBuilder(profileShardResults),
             sortedTopDocs,
             sortValueFormats,
             queryPhaseRankCoordinatorContext,