Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Faster and shorter SearchPhaseController.reduceQueryPhase #119855

Merged
merged 1 commit into from
Jan 9, 2025
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Faster and shorter SearchPhaseController.reduceQueryPhase
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.
  • Loading branch information
original-brownbear committed Jan 9, 2025
commit 6fc4a77fa1eb9dd864d0f3f55f64d8bbe8b507ac
Original file line number Diff line number Diff line change
Expand Up @@ -551,9 +551,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,
Expand All @@ -570,8 +569,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) {
Expand All @@ -581,26 +579,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);
Expand All @@ -611,62 +607,56 @@ 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 InternalAggregations aggregations = bufferedAggs == null
? null
: InternalAggregations.topLevelReduceDelayable(
bufferedAggs,
performFinalReduce ? aggReduceContextBuilder.forFinalReduction() : aggReduceContextBuilder.forPartialReduction()
);
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,
topDocsStats.terminatedEarly,
reducedSuggest,
aggregations,
profileBuilder,
bufferedAggs == null
? null
: InternalAggregations.topLevelReduceDelayable(
bufferedAggs,
performFinalReduce ? aggReduceContextBuilder.forFinalReduction() : aggReduceContextBuilder.forPartialReduction()
),
profileShardResults.isEmpty() ? null : new SearchProfileResultsBuilder(profileShardResults),
sortedTopDocs,
sortValueFormats,
queryPhaseRankCoordinatorContext,
Expand Down